[SandboxVec][BottomUpVec] Add -sbvec-stop-bndl flag for debugging

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

@@ -98,6 +98,7 @@ enum class ResultReason {
   CantSchedule,
   Unimplemented,
   Infeasible,
+  ForcePackForDebugging,
 };
 
 #ifndef NDEBUG
@@ -142,6 +143,8 @@ struct ToStr {
       return "Unimplemented";
     case ResultReason::Infeasible:
       return "Infeasible";
+    case ResultReason::ForcePackForDebugging:
+      return "ForcePackForDebugging";
     }
     llvm_unreachable("Unknown ResultReason enum");
   }
@@ -347,6 +350,10 @@ class LegalityAnalysis {
   // TODO: Try to remove the SkipScheduling argument by refactoring the tests.
   const LegalityResult &canVectorize(ArrayRef<Value *> Bndl,
                                      bool SkipScheduling = false);
+  /// \Returns a Pack with reason 'ForcePackForDebugging'.
+  const LegalityResult &getForcedPackForDebugging() {
+    return createLegalityResult<Pack>(ResultReason::ForcePackForDebugging);
+  }
   void clear();
 };
 

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h

@@ -81,6 +81,10 @@ class BottomUpVec final : public RegionPass {
 #endif // NDEBUG
   };
   ActionsVector Actions;
+  /// Helper counter for debugging. It counts the bundles that we attempt to
+  /// vectorize in vectorizeRec().
+  unsigned DebugBndlCnt = 0;
+
   /// Recursively try to vectorize \p Bndl and its operands. This populates the
   /// `Actions` vector.
   Action *vectorizeRec(ArrayRef<Value *> Bndl, ArrayRef<Value *> UserBndl,

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

@@ -32,6 +32,12 @@ static cl::opt<unsigned long>
            cl::desc("Vectorize if the invocation count is < than this. 0 "
                     "disables vectorization."));
 
+static constexpr const unsigned long StopBundleDisabled =
+    std::numeric_limits<unsigned long>::max();
+static cl::opt<unsigned long>
+    StopBundle("sbvec-stop-bndl", cl::init(StopBundleDisabled), cl::Hidden,
+               cl::desc("Vectorize up to this many bundles."));
+
 namespace sandboxir {
 
 static SmallVector<Value *, 4> getOperand(ArrayRef<Value *> Bndl,
@@ -270,7 +276,10 @@ void BottomUpVec::collectPotentiallyDeadInstrs(ArrayRef<Value *> Bndl) {
 
 Action *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl,
                                   ArrayRef<Value *> UserBndl, unsigned Depth) {
-  const auto &LegalityRes = Legality->canVectorize(Bndl);
+  bool StopForDebug =
+      DebugBndlCnt++ >= StopBundle && StopBundle != StopBundleDisabled;
+  const auto &LegalityRes = StopForDebug ? Legality->getForcedPackForDebugging()
+                                         : Legality->canVectorize(Bndl);
   auto ActionPtr =
       std::make_unique<Action>(&LegalityRes, Bndl, UserBndl, Depth);
   SmallVector<Action *> Operands;
@@ -469,6 +478,7 @@ bool BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl) {
   DeadInstrCandidates.clear();
   Legality->clear();
   Actions.clear();
+  DebugBndlCnt = 0;
   vectorizeRec(Bndl, {}, /*Depth=*/0);
   emitVectors();
   tryEraseDeadInstrs();

llvm/test/Transforms/SandboxVectorizer/stop_bndl.ll

@@ -0,0 +1,71 @@
+; 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="seed-collection<tr-save,bottom-up-vec,tr-accept>" -sbvec-stop-bndl=0 %s -S | FileCheck %s --check-prefix=STOP0
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" -sbvec-stop-bndl=1 %s -S | FileCheck %s --check-prefix=STOP1
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" -sbvec-stop-bndl=2 %s -S | FileCheck %s --check-prefix=STOP2
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" -sbvec-stop-bndl=3 %s -S | FileCheck %s --check-prefix=STOP3
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>"  %s -S | FileCheck %s --check-prefix=NOSTOP
+
+define void @stop_bndl(ptr %ptr) {
+; STOP0-LABEL: define void @stop_bndl(
+; STOP0-SAME: ptr [[PTR:%.*]]) {
+; STOP0-NEXT:    [[PTR0:%.*]] = getelementptr float, ptr [[PTR]], i32 0
+; STOP0-NEXT:    [[PTR1:%.*]] = getelementptr float, ptr [[PTR]], i32 1
+; STOP0-NEXT:    [[LD0:%.*]] = load float, ptr [[PTR0]], align 4
+; STOP0-NEXT:    [[LD1:%.*]] = load float, ptr [[PTR1]], align 4
+; STOP0-NEXT:    [[ADD0:%.*]] = fadd float [[LD0]], 0.000000e+00
+; STOP0-NEXT:    [[ADD1:%.*]] = fadd float [[LD1]], 0.000000e+00
+; STOP0-NEXT:    store float [[ADD0]], ptr [[PTR0]], align 4
+; STOP0-NEXT:    store float [[ADD1]], ptr [[PTR1]], align 4
+; STOP0-NEXT:    ret void
+;
+; STOP1-LABEL: define void @stop_bndl(
+; STOP1-SAME: ptr [[PTR:%.*]]) {
+; STOP1-NEXT:    [[PTR0:%.*]] = getelementptr float, ptr [[PTR]], i32 0
+; STOP1-NEXT:    [[PTR1:%.*]] = getelementptr float, ptr [[PTR]], i32 1
+; STOP1-NEXT:    [[LD0:%.*]] = load float, ptr [[PTR0]], align 4
+; STOP1-NEXT:    [[LD1:%.*]] = load float, ptr [[PTR1]], align 4
+; STOP1-NEXT:    [[ADD0:%.*]] = fadd float [[LD0]], 0.000000e+00
+; STOP1-NEXT:    [[ADD1:%.*]] = fadd float [[LD1]], 0.000000e+00
+; STOP1-NEXT:    [[PACK:%.*]] = insertelement <2 x float> poison, float [[ADD0]], i32 0
+; STOP1-NEXT:    [[PACK1:%.*]] = insertelement <2 x float> [[PACK]], float [[ADD1]], i32 1
+; STOP1-NEXT:    store <2 x float> [[PACK1]], ptr [[PTR0]], align 4
+; STOP1-NEXT:    ret void
+;
+; STOP2-LABEL: define void @stop_bndl(
+; STOP2-SAME: ptr [[PTR:%.*]]) {
+; STOP2-NEXT:    [[PTR0:%.*]] = getelementptr float, ptr [[PTR]], i32 0
+; STOP2-NEXT:    [[PTR1:%.*]] = getelementptr float, ptr [[PTR]], i32 1
+; STOP2-NEXT:    [[LD0:%.*]] = load float, ptr [[PTR0]], align 4
+; STOP2-NEXT:    [[LD1:%.*]] = load float, ptr [[PTR1]], align 4
+; STOP2-NEXT:    [[PACK:%.*]] = insertelement <2 x float> poison, float [[LD0]], i32 0
+; STOP2-NEXT:    [[PACK1:%.*]] = insertelement <2 x float> [[PACK]], float [[LD1]], i32 1
+; STOP2-NEXT:    [[VEC:%.*]] = fadd <2 x float> [[PACK1]], zeroinitializer
+; STOP2-NEXT:    store <2 x float> [[VEC]], ptr [[PTR0]], align 4
+; STOP2-NEXT:    ret void
+;
+; STOP3-LABEL: define void @stop_bndl(
+; STOP3-SAME: ptr [[PTR:%.*]]) {
+; STOP3-NEXT:    [[PTR0:%.*]] = getelementptr float, ptr [[PTR]], i32 0
+; STOP3-NEXT:    [[VECL:%.*]] = load <2 x float>, ptr [[PTR0]], align 4
+; STOP3-NEXT:    [[VEC:%.*]] = fadd <2 x float> [[VECL]], zeroinitializer
+; STOP3-NEXT:    store <2 x float> [[VEC]], ptr [[PTR0]], align 4
+; STOP3-NEXT:    ret void
+;
+; NOSTOP-LABEL: define void @stop_bndl(
+; NOSTOP-SAME: ptr [[PTR:%.*]]) {
+; NOSTOP-NEXT:    [[PTR0:%.*]] = getelementptr float, ptr [[PTR]], i32 0
+; NOSTOP-NEXT:    [[VECL:%.*]] = load <2 x float>, ptr [[PTR0]], align 4
+; NOSTOP-NEXT:    [[VEC:%.*]] = fadd <2 x float> [[VECL]], zeroinitializer
+; NOSTOP-NEXT:    store <2 x float> [[VEC]], ptr [[PTR0]], align 4
+; NOSTOP-NEXT:    ret void
+;
+  %ptr0 = getelementptr float, ptr %ptr, i32 0
+  %ptr1 = getelementptr float, ptr %ptr, i32 1
+  %ld0 = load float, ptr %ptr0
+  %ld1 = load float, ptr %ptr1
+  %add0 = fadd float %ld0, 0.0
+  %add1 = fadd float %ld1, 0.0
+  store float %add0, ptr %ptr0
+  store float %add1, ptr %ptr1
+  ret void
+}