[llvm] [RISCV] Combine build_vector of extract_vector_elts to vector_shuffle (PR #80883)

[Philip Reames via llvm-commits]

================
@@ -14933,11 +14933,68 @@ static SDValue performSELECTCombine(SDNode *N, SelectionDAG &DAG,
   return tryFoldSelectIntoOp(N, DAG, FalseVal, TrueVal, /*Swapped*/true);
 }
 
-/// If we have a build_vector where each lane is binop X, C, where C
-/// is a constant (but not necessarily the same constant on all lanes),
-/// form binop (build_vector x1, x2, ...), (build_vector c1, c2, c3, ..).
-/// We assume that materializing a constant build vector will be no more
-/// expensive that performing O(n) binops.
+/// Canonicalize
+///
+/// (build_vector (extract_vector_elt v, i0),
+///               (extract_vector_elt v, i1),
+///               (extract_vector_elt v, i2),
+///               (extract_vector_elt v, i3))
+///
+/// to
+///
+/// (vector_shuffle<i0, i1, i2, i3> v, undef)
+///
+/// shufflevectors may be lowered to the build_vector pattern above if the
+/// vector types don't match, so try and recover the shuffle to avoid
+/// scalarization.
+///
+/// Note that this combine exists in a similar form in DAGCombiner.cpp, but it
+/// depends on extract_subvector being cheap. On RISC-V we always want to use
+/// the vector_shuffle form, so we perform it unconditionally here.
+static SDValue combineBuildVectorToShuffle(SDNode *N, SelectionDAG &DAG,
+                                           const RISCVTargetLowering &TLI) {
+  EVT VT = N->getValueType(0);
+  SDLoc DL(N);
+  if (!TLI.isTypeLegal(VT))
+    return SDValue();
+
+  SDValue ExtractVec;
+  SmallVector<int> Mask(N->getNumOperands(), -1);
+  for (auto [i, Op] : enumerate(N->op_values())) {
+    if (Op.isUndef())
+      continue;
+    if (Op.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
+      return SDValue();
+    if (!isa<ConstantSDNode>(Op.getOperand(1)))
+      return SDValue();
+    if (!ExtractVec) {
+      ExtractVec = Op.getOperand(0);
+      Mask[i] = Op.getConstantOperandVal(1);
+      continue;
+    }
+    if (Op.getOperand(0) != ExtractVec)
+      return SDValue();
+    Mask[i] = Op.getConstantOperandVal(1);
+  }
+
+  EVT ExtractVT = ExtractVec.getValueType();
+  if (ExtractVT.isScalableVector())
----------------
preames wrote:

You probably should check that ExtractVT is legal here.

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