[llvm] [AArch64] Extend performActiveLaneMaskCombine for more than two extracts (PR #146725)

[Kerry McLaughlin via llvm-commits]

================
@@ -18143,53 +18143,63 @@ performActiveLaneMaskCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
       (!ST->hasSVE2p1() && !(ST->hasSME2() && ST->isStreaming())))
     return SDValue();
 
-  if (!N->hasNUsesOfValue(2, 0))
+  unsigned NumUses = N->use_size();
+  unsigned MaskMinElts = N->getValueType(0).getVectorMinNumElements();
+  if (MaskMinElts % NumUses != 0)
     return SDValue();
 
-  const uint64_t HalfSize = N->getValueType(0).getVectorMinNumElements() / 2;
-  if (HalfSize < 2)
+  unsigned ExtMinElts = MaskMinElts / NumUses;
+  if (ExtMinElts < 2)
     return SDValue();
 
-  auto It = N->user_begin();
-  SDNode *Lo = *It++;
-  SDNode *Hi = *It;
+  SmallVector<SDNode *> Extracts(NumUses, nullptr);
+  for (SDNode *Use : N->users()) {
+    if (Use->getOpcode() != ISD::EXTRACT_SUBVECTOR)
+      return SDValue();
 
-  if (Lo->getOpcode() != ISD::EXTRACT_SUBVECTOR ||
-      Hi->getOpcode() != ISD::EXTRACT_SUBVECTOR)
-    return SDValue();
+    // Ensure the extract type is correct (e.g. if NumUses is 4 and
+    // the mask return type is nxv8i1, each extract should be nxv2i1.
+    if (Use->getValueType(0).getVectorMinNumElements() != ExtMinElts)
----------------
kmclaughlin-arm wrote:

I think this should be covered by `!N->getValueType(0).isScalableVector()` at the start of the function, and there is an existing test (`@test_fixed_extract`) which is testing fixed extracts from a scalable mask.

The test is extracting two `<2 x i1>` from one `<vscale x 8 x i1>` mask however, so I have changed it slightly so that the only reason it should be rejected is because the return types of the extracts are not scalable.

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