[llvm-branch-commits] [llvm] [LoongArch] Custom legalize vector_shuffle to `xvinsve0.{w/d}` when possible (PR #160857)

[via llvm-branch-commits]

================
@@ -2317,6 +2317,54 @@ static SDValue lowerVECTOR_SHUFFLE_XVPICKOD(const SDLoc &DL, ArrayRef<int> Mask,
   return DAG.getNode(LoongArchISD::VPICKOD, DL, VT, V2, V1);
 }
 
+// Check if exactly one element of the Mask is replaced by 'Replaced', while
+// all other elements are either 'Base + i' or undef (-1). On success, return
+// the index of the replaced element. Otherwise, just return -1.
+static int checkReplaceOne(ArrayRef<int> Mask, int Base, int Replaced) {
+  int MaskSize = Mask.size();
+  int Idx = -1;
+  for (int i = 0; i < MaskSize; ++i) {
+    if (Mask[i] == Base + i || Mask[i] == -1)
+      continue;
+    if (Mask[i] != Replaced)
+      return -1;
+    if (Idx == -1)
+      Idx = i;
+    else
+      return -1;
+  }
+  return Idx;
+}
+
+/// Lower VECTOR_SHUFFLE into XVINSVE0 (if possible).
+static SDValue
+lowerVECTOR_SHUFFLE_XVINSVE0(const SDLoc &DL, ArrayRef<int> Mask, MVT VT,
+                             SDValue V1, SDValue V2, SelectionDAG &DAG,
+                             const LoongArchSubtarget &Subtarget) {
+  // LoongArch LASX only supports xvinsve0.{w/d}.
+  if (VT != MVT::v8i32 && VT != MVT::v8f32 && VT != MVT::v4i64 &&
+      VT != MVT::v4f64)
+    return SDValue();
+
+  MVT GRLenVT = Subtarget.getGRLenVT();
+  int MaskSize = Mask.size();
+  assert(MaskSize == (int)VT.getVectorNumElements() && "Unexpected mask size");
+
+  // Case 1: the lowest element of V2 replaces one element in V1.
+  int Idx = checkReplaceOne(Mask, 0, MaskSize);
+  if (Idx != -1)
+    return DAG.getNode(LoongArchISD::XVINSVE0, DL, VT, V1, V2,
+                       DAG.getConstant(Idx, DL, GRLenVT));
+
+  // Case 2: the lowest element of V1 replaces one element in V2.
+  Idx = checkReplaceOne(Mask, MaskSize, 0);
----------------
zhaoqi5 wrote:

Thanks for your great analysis!

The new version has more complex `for` loop than the old version because both cases should be checked at meantime. In the old version, it could return early in some cases, and would not enter the twice interation if the first case successed.

I have not come up with a better solution. I will revert to the old version (and maybe also use a lambda function).

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