[llvm] [SelectionDAG] Fold extracts spanning concat operands (PR #200936)

[Paul Walker via llvm-commits]

================
@@ -27537,6 +27537,64 @@ static SDValue foldExtractSubvectorFromShuffleVector(EVT NarrowVT, SDValue Src,
   return DAG.getVectorShuffle(NarrowVT, DL, NewOps[0], NewOps[1], NewMask);
 }
 
+static SDValue foldExtractSubvectorFromConcatVectors(EVT VT, SDValue V,
+                                                     uint64_t ExtIdx,
+                                                     const SDLoc &DL,
+                                                     SelectionDAG &DAG,
+                                                     bool LegalOperations) {
+  if (V.getOpcode() != ISD::CONCAT_VECTORS)
+    return SDValue();
+
+  ElementCount ExtNumElts = VT.getVectorElementCount();
+  EVT ConcatSrcVT = V.getOperand(0).getValueType();
+
+  ElementCount ConcatSrcNumElts = ConcatSrcVT.getVectorElementCount();
+  unsigned ConcatOpIdx = ExtIdx / ConcatSrcNumElts.getKnownMinValue();
+  if (ConcatOpIdx >= V.getNumOperands())
+    return SDValue();
+
+  // If the concatenated source types match this extract, it's a direct
+  // simplification:
+  // extract_subvec (concat V1, V2, ...), i --> Vi
+  if (VT.getVectorElementCount() == ConcatSrcVT.getVectorElementCount())
+    return V.getOperand(ConcatOpIdx);
+
+  // If the concatenated source vectors are a multiple length of this extract,
+  // then extract a fraction of one of those source vectors directly from a
+  // concat operand. Example:
+  //   v2i8 extract_subvec (v16i8 concat (v8i8 X), (v8i8 Y)), 14 -->
+  //   v2i8 extract_subvec v8i8 Y, 6
+  if (ConcatSrcNumElts.isKnownMultipleOf(ExtNumElts) &&
+      ConcatSrcNumElts.isFixed() == ExtNumElts.isFixed()) {
+    uint64_t NewExtIdx =
+        ExtIdx - ConcatOpIdx * ConcatSrcNumElts.getKnownMinValue();
+    if (NewExtIdx + ExtNumElts.getKnownMinValue() >
+        ConcatSrcNumElts.getKnownMinValue())
+      return SDValue();
+    return DAG.getExtractSubvector(DL, VT, V.getOperand(ConcatOpIdx),
+                                   NewExtIdx);
+  }
+
+  // If the extract covers multiple whole concat operands, rebuild that smaller
+  // concat directly.
+  bool IsPermittedConcat =
+      !LegalOperations || DAG.getTargetLoweringInfo().isOperationLegalOrCustom(
+                              ISD::CONCAT_VECTORS, VT);
+  if (ExtNumElts.isKnownMultipleOf(ConcatSrcNumElts) &&
+      ExtIdx % ConcatSrcNumElts.getKnownMinValue() == 0 &&
+      ExtNumElts.isFixed() == ConcatSrcNumElts.isFixed() && IsPermittedConcat) {
+    unsigned NumConcatOps =
+        ExtNumElts.getKnownMinValue() / ConcatSrcNumElts.getKnownMinValue();
----------------
paulwalker-arm wrote:

```suggestion
    unsigned NumConcatOps = ExtNumElts.getKnownScalarFactor(ConcatSrcNumElts);
```

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