[llvm] [AArch64] Add DAG combine to widen v3i8/v4i8 VECTOR_MATCH needles (PR #218972)

[Paul Walker via llvm-commits]

================
@@ -1656,6 +1656,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
 
       for (MVT VT : {MVT::v16i1, MVT::v8i1, MVT::v16i8, MVT::v8i8})
         setOperationAction(ISD::VECTOR_MATCH, VT, Custom);
+
+      setTargetDAGCombine(ISD::VECTOR_MATCH);
----------------
paulwalker-arm wrote:

You are overcomplicating this. `setOperationAction()` does not apply any special meaning to the specified type. It is the users of `getTypeAction()` that do that when deciding where to pull the type from. While not all actions are applicable to both type and operation legalisation, `Custom` is an exception specifically for this use case.

It would have been nice if you tried the suggestion before dismissing it but for the avoidance of doubt the following matches this PR's output:
```
@@ -1654,10 +1654,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       for (MVT VT : {MVT::nxv16i1, MVT::nxv8i1})
         setOperationAction(ISD::VECTOR_MATCH, VT, Custom);
 
-      for (MVT VT : {MVT::v16i1, MVT::v8i1, MVT::v16i8, MVT::v8i8})
+      for (MVT VT : {MVT::v16i1, MVT::v8i1, MVT::v16i8, MVT::v8i8, MVT::v3i8, MVT::v4i8})
         setOperationAction(ISD::VECTOR_MATCH, VT, Custom);
-
-      setTargetDAGCombine(ISD::VECTOR_MATCH);
     }
 
     setOperationAction(ISD::GET_ACTIVE_LANE_MASK, MVT::nxv1i1, Custom);
@@ -6519,6 +6517,31 @@ static SDValue LowerVectorMatch(SDValue Op, SelectionDAG &DAG) {
           Op1VT.getVectorElementType() == MVT::i16) &&
          "Expected 8-bit or 16-bit characters.");
 
+  if ((Op2VT == MVT::v3i8 || Op2VT == MVT::v4i8)) {
+    SDValue Needle = Op2;
+    EVT NeedleVT = Op2VT;
+
+    if (NeedleVT == MVT::v3i8) {
+      // Pad a v3i8 needle to v4i8.
+      SDValue Pad = DAG.getExtractVectorElt(DL, MVT::i8, Needle, 0);
+      Needle = DAG.getInsertSubvector(DL, DAG.getPOISON(MVT::v4i8), Needle, 0);
+      Needle = DAG.getInsertVectorElt(DL, Needle, Pad, 3);
+    }
+
+    // Splat the needle to a full scalable vector (in such a way that the bitcasts
+    // and extracts should fold away).
+    Needle = DAG.getBitcast(MVT::v1i32, Needle);
+    Needle = DAG.getExtractVectorElt(DL, MVT::i32, Needle, 0);
+    Needle = DAG.getSplatVector(MVT::nxv4i32, DL, Needle);
+    Needle = DAG.getBitcast(MVT::nxv16i8, Needle);
+    Needle = DAG.getExtractSubvector(DL, MVT::v16i8, Needle, 0);
+
+    return DAG.getNode(ISD::VECTOR_MATCH, DL, Op.getValueType(),
+        Op1, Needle, Mask);
+  }
```
Because the setOperationAction change might try custom type legalisation for matching result types you might need to ensure ReplaceNodeResults ignore them

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