[SandboxVec][Scheduler] Don't allow rescheduling of already scheduled (#128050)

This patch implements the check for not allowing re-scheduling of
instructions that have already been scheduled in a scheduling bundle.
Rescheduling should only happen if the instructions were temporarily
scheduled in singleton bundles during a previous call to
`trySchedule()`.

Author: vporpo

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h

@@ -133,6 +133,10 @@ class SchedBundle {
       N->clearSchedBundle();
   }
   bool empty() const { return Nodes.empty(); }
+  /// Singleton bundles are created when scheduling instructions temporarily to
+  /// fill in the schedule until we schedule the vector bundle. These are
+  /// non-vector bundles containing just a single instruction.
+  bool isSingleton() const { return Nodes.size() == 1u; }
   DGNode *back() const { return Nodes.back(); }
   using iterator = ContainerTy::iterator;
   using const_iterator = ContainerTy::const_iterator;
@@ -190,10 +194,12 @@ class Scheduler {
   /// The scheduling state of the instructions in the bundle.
   enum class BndlSchedState {
     NoneScheduled, ///> No instruction in the bundle was previously scheduled.
-    PartiallyOrDifferentlyScheduled, ///> Only some of the instrs in the bundle
-                                     /// were previously scheduled, or all of
-                                     /// them were but not in the same
-                                     /// SchedBundle.
+    AlreadyScheduled, ///> At least one instruction in the bundle belongs to a
+                      /// different non-singleton scheduling bundle.
+    TemporarilyScheduled, ///> Instructions were temporarily scheduled as
+                          /// singleton bundles or some of them were not
+                          /// scheduled at all. None of them were in a vector
+                          ///(non-singleton) bundle.
     FullyScheduled, ///> All instrs in the bundle were previously scheduled and
                     /// were in the same SchedBundle.
   };
@@ -246,6 +252,11 @@ class Scheduler {
 class SchedulerInternalsAttorney {
 public:
   static DependencyGraph &getDAG(Scheduler &Sched) { return Sched.DAG; }
+  using BndlSchedState = Scheduler::BndlSchedState;
+  static BndlSchedState getBndlSchedState(const Scheduler &Sched,
+                                          ArrayRef<Instruction *> Instrs) {
+    return Sched.getBndlSchedState(Instrs);
+  }
 };
 
 } // namespace llvm::sandboxir

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

@@ -161,30 +161,36 @@ bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
 Scheduler::BndlSchedState
 Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
   assert(!Instrs.empty() && "Expected non-empty bundle");
-  bool PartiallyScheduled = false;
-  bool FullyScheduled = true;
-  for (auto *I : Instrs) {
+  auto *N0 = DAG.getNode(Instrs[0]);
+  auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr;
+  bool AllUnscheduled = SB0 == nullptr;
+  bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton();
+  for (auto *I : drop_begin(Instrs)) {
     auto *N = DAG.getNode(I);
-    if (N != nullptr && N->scheduled())
-      PartiallyScheduled = true;
-    else
-      FullyScheduled = false;
-  }
-  if (FullyScheduled) {
-    // If not all instrs in the bundle are in the same SchedBundle then this
-    // should be considered as partially-scheduled, because we will need to
-    // re-schedule.
-    SchedBundle *SB = DAG.getNode(Instrs[0])->getSchedBundle();
-    assert(SB != nullptr && "FullyScheduled assumes that there is an SB!");
-    if (any_of(drop_begin(Instrs), [this, SB](sandboxir::Value *SBV) {
-          return DAG.getNode(cast<sandboxir::Instruction>(SBV))
-                     ->getSchedBundle() != SB;
-        }))
+    auto *SB = N != nullptr ? N->getSchedBundle() : nullptr;
+    if (SB != nullptr) {
+      // We found a scheduled instr, so there is now way all are unscheduled.
+      AllUnscheduled = false;
+      if (SB->isSingleton()) {
+        // We found an instruction in a temporarily scheduled singleton. There
+        // is no way that all instructions are scheduled in the same bundle.
+        FullyScheduled = false;
+      }
+    }
+
+    if (SB != SB0) {
+      // Either one of SB, SB0 is null, or they are in different bundles, so
+      // Instrs are definitely not in the same vector bundle.
       FullyScheduled = false;
+      // One of SB, SB0 are in a vector bundle and they differ.
+      if ((SB != nullptr && !SB->isSingleton()) ||
+          (SB0 != nullptr && !SB0->isSingleton()))
+        return BndlSchedState::AlreadyScheduled;
+    }
   }
-  return FullyScheduled       ? BndlSchedState::FullyScheduled
-         : PartiallyScheduled ? BndlSchedState::PartiallyOrDifferentlyScheduled
-                              : BndlSchedState::NoneScheduled;
+  return AllUnscheduled   ? BndlSchedState::NoneScheduled
+         : FullyScheduled ? BndlSchedState::FullyScheduled
+                          : BndlSchedState::TemporarilyScheduled;
 }
 
 void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
@@ -203,13 +209,14 @@ void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
   //
   Instruction *TopI = &*ScheduleTopItOpt.value();
   Instruction *LowestI = VecUtils::getLowest(Instrs);
-  // Destroy the schedule bundles from LowestI all the way to the top.
+  // Destroy the singleton schedule bundles from LowestI all the way to the top.
   for (auto *I = LowestI, *E = TopI->getPrevNode(); I != E;
        I = I->getPrevNode()) {
     auto *N = DAG.getNode(I);
     if (N == nullptr)
       continue;
-    if (auto *SB = N->getSchedBundle())
+    auto *SB = N->getSchedBundle();
+    if (SB->isSingleton())
       eraseBundle(SB);
   }
   // The DAG Nodes contain state like the number of UnscheduledSuccs and the
@@ -259,7 +266,12 @@ bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
   case BndlSchedState::FullyScheduled:
     // Nothing to do.
     return true;
-  case BndlSchedState::PartiallyOrDifferentlyScheduled:
+  case BndlSchedState::AlreadyScheduled:
+    // Instructions are part of a different vector schedule, so we can't
+    // schedule \p Instrs in the same bundle (without destroying the existing
+    // schedule).
+    return false;
+  case BndlSchedState::TemporarilyScheduled:
     // If one or more instrs are already scheduled we need to destroy the
     // top-most part of the schedule that includes the instrs in the bundle and
     // re-schedule.

llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll

@@ -386,3 +386,33 @@ define void @vecInstrsPlacement(ptr %ptr0) {
   store double %add1, ptr %ptr1
   ret void
 }
+
+; During the bottom-up traversal we form bundle {ldA0,ldA1} but later when we
+; visit the RHS operands of the additions we try to form {ldA1,ldA2}
+; which is not allowed.
+define void @instrsInMultipleBundles(ptr noalias %ptr) {
+; CHECK-LABEL: define void @instrsInMultipleBundles(
+; CHECK-SAME: ptr noalias [[PTR:%.*]]) {
+; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr i8, ptr [[PTR]], i64 0
+; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr i8, ptr [[PTR]], i64 2
+; CHECK-NEXT:    [[LDA2:%.*]] = load i8, ptr [[GEP2]], align 1
+; CHECK-NEXT:    [[VECL:%.*]] = load <2 x i8>, ptr [[GEP0]], align 1
+; CHECK-NEXT:    [[VEXT:%.*]] = extractelement <2 x i8> [[VECL]], i32 1
+; CHECK-NEXT:    [[VINS:%.*]] = insertelement <2 x i8> poison, i8 [[VEXT]], i32 0
+; CHECK-NEXT:    [[VINS1:%.*]] = insertelement <2 x i8> [[VINS]], i8 [[LDA2]], i32 1
+; CHECK-NEXT:    [[VEC:%.*]] = add <2 x i8> [[VECL]], [[VINS1]]
+; CHECK-NEXT:    store <2 x i8> [[VEC]], ptr [[GEP0]], align 1
+; CHECK-NEXT:    ret void
+;
+  %gep0 = getelementptr i8, ptr %ptr, i64 0
+  %gep1 = getelementptr i8, ptr %ptr, i64 1
+  %gep2 = getelementptr i8, ptr %ptr, i64 2
+  %ldA0 = load i8, ptr %gep0
+  %ldA1 = load i8, ptr %gep1
+  %ldA2 = load i8, ptr %gep2
+  %add0 = add i8 %ldA0, %ldA1
+  %add1 = add i8 %ldA1, %ldA2
+  store i8 %add0, ptr %gep0
+  store i8 %add1, ptr %gep1
+  ret void
+}