[llvm] [SLP] NFC. Add assert for ShuffleCostEstimator and ShuffleInstructionBuilder usage. (PR #123708)

[Han-Kuan Chen via llvm-commits]

https://github.com/HanKuanChen created https://github.com/llvm/llvm-project/pull/123708

None

>From 6b87611ff97f3f8d057539df28d191c975d4cf0c Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Thu, 16 Jan 2025 04:50:40 -0800
Subject: [PATCH] [SLP] NFC. Add assert for ShuffleCostEstimator and
 ShuffleInstructionBuilder usage.

---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index fc3afed391a06d..5140a8ec61654e 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10932,6 +10932,7 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
            ArrayRef<std::pair<const TreeEntry *, unsigned>> SubVectors,
            ArrayRef<int> SubVectorsMask, unsigned VF = 0,
            function_ref<void(Value *&, SmallVectorImpl<int> &)> Action = {}) {
+    assert(!IsFinalized && "ShuffleCostEstimator is already finalized.");
     IsFinalized = true;
     if (Action) {
       const PointerUnion<Value *, const TreeEntry *> &Vec = InVectors.front();
@@ -14524,6 +14525,7 @@ class BoUpSLP::ShuffleInstructionBuilder final : public BaseShuffleAnalysis {
            ArrayRef<std::pair<const TreeEntry *, unsigned>> SubVectors,
            ArrayRef<int> SubVectorsMask, unsigned VF = 0,
            function_ref<void(Value *&, SmallVectorImpl<int> &)> Action = {}) {
+    assert(!IsFinalized && "ShuffleInstructionBuilder is already finalized.");
     IsFinalized = true;
     unsigned ScalarTyNumElements = getNumElements(ScalarTy);
     SmallVector<int> NewExtMask(ExtMask);
@@ -14696,6 +14698,8 @@ Value *BoUpSLP::vectorizeOperand(TreeEntry *E, unsigned NodeIdx,
           NumElements != 1 ? FixedVectorType::get(ScalarTy, NumElements)
                            : ScalarTy,
           Builder, *this);
+      assert(Mask.size() == VL.size() &&
+             "The mask is incompatible with the expected result.");
       ShuffleBuilder.add(V, Mask);
       SmallVector<std::pair<const TreeEntry *, unsigned>> SubVectors(
           E->CombinedEntriesWithIndices.size());
@@ -14912,6 +14916,8 @@ ResTy BoUpSLP::processBuildVector(const TreeEntry *E, Type *ScalarTy,
         // process to keep correct order.
         return *Delayed;
       }
+      assert(ExtractMask.size() == E->Scalars.size() &&
+             "The mask is incompatible with the expected result.");
       if (Value *VecBase = ShuffleBuilder.adjustExtracts(
               E, ExtractMask, ExtractShuffles, NumParts, UseVecBaseAsInput)) {
         ExtractVecBase = VecBase;
@@ -15229,6 +15235,9 @@ ResTy BoUpSLP::processBuildVector(const TreeEntry *E, Type *ScalarTy,
       SmallVector<int> BVMask(GatheredScalars.size(), PoisonMaskElem);
       TryPackScalars(GatheredScalars, BVMask, /*IsRootPoison=*/true);
       Value *BV = ShuffleBuilder.gather(GatheredScalars, BVMask.size());
+      assert((BVMask.size() == E->Scalars.size() ||
+              BVMask.size() == E->ReuseShuffleIndices.size()) &&
+             "The mask is incompatible with the expected result.");
       ShuffleBuilder.add(BV, BVMask);
     }
     if (all_of(NonConstants, [=](Value *V) {
@@ -15250,6 +15259,9 @@ ResTy BoUpSLP::processBuildVector(const TreeEntry *E, Type *ScalarTy,
     SmallVector<int> ReuseMask(GatheredScalars.size(), PoisonMaskElem);
     TryPackScalars(GatheredScalars, ReuseMask, /*IsRootPoison=*/true);
     Value *BV = ShuffleBuilder.gather(GatheredScalars, ReuseMask.size());
+    assert((ReuseMask.size() == E->Scalars.size() ||
+            ReuseMask.size() == E->ReuseShuffleIndices.size()) &&
+           "The mask is incompatible with the expected result.");
     ShuffleBuilder.add(BV, ReuseMask);
     Res = ShuffleBuilder.finalize(E->ReuseShuffleIndices, SubVectors,
                                   SubVectorsMask);
@@ -15261,6 +15273,8 @@ ResTy BoUpSLP::processBuildVector(const TreeEntry *E, Type *ScalarTy,
         Mask[I] = I;
     }
     Value *BV = ShuffleBuilder.gather(GatheredScalars);
+    assert(Mask.size() == E->Scalars.size() &&
+           "The mask is incompatible with the expected result.");
     ShuffleBuilder.add(BV, Mask);
     Res = ShuffleBuilder.finalize(E->ReuseShuffleIndices, SubVectors,
                                   SubVectorsMask);
@@ -15329,10 +15343,15 @@ Value *BoUpSLP::vectorizeTree(TreeEntry *E, bool PostponedPHIs) {
       ArrayRef<int> Mask =
           ArrayRef(reinterpret_cast<const int *>(E->ReorderIndices.begin()),
                    E->ReorderIndices.size());
+      assert((Mask.empty() || Mask.size() == E->Scalars.size()) &&
+             "The mask is incompatible with the expected result.");
       ShuffleBuilder.add(V, Mask);
     } else if (E->State == TreeEntry::StridedVectorize && IsReverseOrder) {
       ShuffleBuilder.addOrdered(V, {});
     } else {
+      assert((E->ReorderIndices.empty() ||
+              E->ReorderIndices.size() == E->Scalars.size()) &&
+             "The mask is incompatible with the expected result.");
       ShuffleBuilder.addOrdered(V, E->ReorderIndices);
     }
     SmallVector<std::pair<const TreeEntry *, unsigned>> SubVectors(