[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())
----------------
paulwalker-arm wrote:

Can this happen? NewExtIdx is effectively "ExtIdx % ConcatSrcNumElts" and we already know:
```
ConcatSrcNumElts = A * ExtNumElts
ExtIdx = B * ExtNumElts
```
Therefore NewExtIdx must also be a multiple of ExtNumElts and thus cannot overflow?

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