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

[Matt Arsenault 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
----------------
arsenm wrote:

```suggestion
  //   v2i8 extract_subvector (v16i8 concat v8i8:X, v8i8 :), 14 -->
  //   v2i8 extract_subvector v8i8:Y, 6
```

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