[llvm] [VectorCombine] Fold compare chains to experimental.vector.match (PR #212456)

[Graham Hunter via llvm-commits]

================
@@ -578,6 +579,123 @@ bool VectorCombine::isExtractExtractCheap(ExtractElementInst *Ext0,
   return OldCost < NewCost;
 }
 
+/// Attempts to fold chains of `eq` comparisons to `vector.match`. For example:
+///
+/// ```
+/// %c0 = icmp eq <vscale x 4 x i16> %vec, splat (i16 100)
+/// %c1 = icmp eq <vscale x 4 x i16> %vec, splat (i16 32)
+/// %c2 = icmp eq <vscale x 4 x i16> %vec, splat (i16 10)
+/// %or = or <vscale x 4 x i1> %c0, %c1
+/// %r = or <vscale x 4 x i1> %or, %c2
+/// ```
+/// folds to:
+/// ```
+/// %r = call <vscale x 4 x i1> @llvm.experimental.vector.match.nxv8i16.v8i16(
+///        <vscale x 4 x i16> %vec, <i16 100, i16 32, i16 10, i16 100>,
+///        splat (i1 true))
+/// ```
+bool VectorCombine::foldDisjunctionToConstantMatch(Instruction &I) {
+  auto *BinOp = dyn_cast<BinaryOperator>(&I);
+  if (!BinOp || BinOp->getOpcode() != Instruction::Or)
+    return false;
+
+  // The vector being checked (%vec in the example).
+  Value *CompareSource = nullptr;
+  // The values being searched for (e.g, 100, 32, 10).
+  SmallSetVector<Constant *, 16> SearchValues;
+
+  SmallPtrSet<Value *, 4> Visited;
+  SmallVector<Value *> Worklist = {BinOp->getOperand(0), BinOp->getOperand(1)};
+  unsigned NumVisited = 0;
+  while (!Worklist.empty()) {
+    auto *Op = dyn_cast<Instruction>(Worklist.pop_back_val());
+    if (++NumVisited >= MaxInstrsToScan || !Op || !Op->hasOneUse())
+      return false;
+
+    if (!Visited.insert(Op).second)
+      continue;
+
+    Constant *C;
+    Value *Source = nullptr;
+    // TODO: Extend to "not equals" compares.
+    if (match(Op, m_c_SpecificICmp(CmpInst::ICMP_EQ, m_Value(Source),
+                                   m_ConstantSplat(m_Constant(C)))) &&
+        !C->getType()->isPointerTy()) {
+      SearchValues.insert(C);
+    } else if (match(Op, m_Intrinsic<Intrinsic::experimental_vector_match>(
+                             m_Value(Source), m_Constant(C),
+                             m_ConstantSplat(m_SpecificInt(1))))) {
+      // Merge with any previous vector match. This ensures the combine always
+      // creates the largest possible match (even if the combine applies to an
+      // `or` earlier in the chain).
+      unsigned NumElts = cast<FixedVectorType>(C->getType())->getNumElements();
+      for (unsigned I = 0; I < NumElts; ++I)
+        if (Constant *Elt = C->getAggregateElement(I))
+          SearchValues.insert(Elt);
+    } else if (Op->getOpcode() == Instruction::Or) {
+      Worklist.append({Op->getOperand(0), Op->getOperand(1)});
+    } else {
+      // InstCombine can fold compares with nearby values to check ranges, which
+      // means this combine does not match. TODO: Extend this fold to cover
+      // these cases too?
+      return false;
+    }
+
+    // Ensure all comparisons are against the same vector.
+    if (!CompareSource)
+      CompareSource = Source;
+    else if (Source && CompareSource != Source)
+      return false;
+  }
+
+  if (SearchValues.size() <= 1)
+    return false;
+
+  auto *ResultTy = cast<VectorType>(BinOp->getType());
+  auto *SrcType = cast<VectorType>(CompareSource->getType());
+
+  // Cost the old code as-if it were all a chain of compares. It could include
+  // a match, but that will always be replaced with another match (with a
+  // needle size >= to the current match).
+  InstructionCost OldCost =
+      TTI.getCmpSelInstrCost(Instruction::ICmp, SrcType, ResultTy,
+                             CmpInst::ICMP_EQ, CostKind) *
+      SearchValues.size();
+  OldCost += TTI.getArithmeticInstrCost(Instruction::Or, ResultTy, CostKind) *
+             (SearchValues.size() - 1);
+
+  ElementCount SrcElts = SrcType->getElementCount();
+
+  InstructionCost NewCost;
+  Constant *NeedleVector;
+  SmallVector<Constant *> MatchValues(SearchValues.getArrayRef());
+
+  // Look for a match needle size (up to the size of SrcType) that's profitable.
+  do {
+    NeedleVector = ConstantVector::get(MatchValues);
+    IntrinsicCostAttributes ICA(Intrinsic::experimental_vector_match, ResultTy,
+                                {SrcType, NeedleVector->getType(), ResultTy});
+    NewCost = TTI.getIntrinsicInstrCost(ICA, CostKind);
+    if (NewCost < OldCost)
+      break;
+
+    // Pad the needle by duplicating the first element.
+    unsigned PadValues = NextPowerOf2(MatchValues.size()) - MatchValues.size();
+    MatchValues.append(SmallVector<Constant *>(PadValues, MatchValues[0]));
----------------
huntergr-arm wrote:

```suggestion
    MatchValues.resize(NextPowerOf2(MatchValues.size()), MatchValues.front());
```

https://github.com/llvm/llvm-project/pull/212456