[llvm] [SDAG][AArch64] Fold extract from pext to use status flags (PR #206443)

[Gaëtan Bossu via llvm-commits]

================
@@ -21970,6 +21970,65 @@ performExtractLastActiveCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
                      Vec);
 }
 
+static bool hasSVE2p1OrStreamingSME2(const AArch64Subtarget *Subtarget) {
+  return (Subtarget->isSVEorStreamingSVEAvailable() &&
+          Subtarget->hasSVE2p1()) ||
+         (Subtarget->isStreaming() && Subtarget->hasSME2());
+}
+
+static auto m_PredicateAsCounterWhile() {
+  using namespace llvm::SDPatternMatch;
+  return m_AnyOf(m_SpecificOpc(AArch64ISD::WHILEGE_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILEGT_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILELT_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILELE_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILEHS_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILEHI_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILELO_PRED_COUNTER),
+                 m_SpecificOpc(AArch64ISD::WHILELS_PRED_COUNTER));
+}
+
+/// Folds extracting the first lane from the first segment of a
+/// predicate-as-counter while to a conditional set (CSET) based on the
+/// "FIRST_ACTIVE" status flag from the while.
+///
+///   %while = WHILE_*_PRED_COUNTER .. ; predicate-as-counter while
+///   %first.segment = pext(%while, 0) ; predicate extract of segment 0
+///   %first.active = extract_elt(%first_segment, 0) ; extract first lane
+///
+///   ->
+///
+///   %while = WHILE_*_PRED_COUNTER .. ; predicate-as-counter while
+///   %first.active = cset(%while, FIRST_ACTIVE)
----------------
gbossu wrote:

Nit:
```
%while, %flags = WHILE_*_PRED_COUNTER
...
->
%while, %flags = WHILE_*_PRED_COUNTER
%first.active = cset(%flags, FIRST_ACTIVE)
```

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