[llvm] [LoongArch] Optimize extractelement containing variable index (PR #151475)

[via llvm-commits]

================
@@ -2608,14 +2613,91 @@ SDValue LoongArchTargetLowering::lowerCONCAT_VECTORS(SDValue Op,
 SDValue
 LoongArchTargetLowering::lowerEXTRACT_VECTOR_ELT(SDValue Op,
                                                  SelectionDAG &DAG) const {
-  EVT VecTy = Op->getOperand(0)->getValueType(0);
+  MVT EltVT = Op.getSimpleValueType();
+  SDValue Vec = Op->getOperand(0);
+  EVT VecTy = Vec->getValueType(0);
   SDValue Idx = Op->getOperand(1);
-  unsigned NumElts = VecTy.getVectorNumElements();
+  SDLoc DL(Op);
+  MVT GRLenVT = Subtarget.getGRLenVT();
+
+  assert(VecTy.is256BitVector() && "Unexpected EXTRACT_VECTOR_ELT vector type");
 
-  if (isa<ConstantSDNode>(Idx) && Idx->getAsZExtVal() < NumElts)
+  if (isa<ConstantSDNode>(Idx))
     return Op;
 
-  return SDValue();
+  switch (VecTy.getSimpleVT().SimpleTy) {
+  default:
+    llvm_unreachable("Unexpected type");
+  case MVT::v32i8:
+  case MVT::v16i16: {
+    // Consider the source vector as v8i32 type.
+    SDValue NewVec = DAG.getBitcast(MVT::v8i32, Vec);
+
+    // Compute the adjusted index and use it to broadcast the vector.
+    // The original desired i8/i16 element is now replicated in each
+    // i32 lane of the splatted vector.
+    SDValue NewIdx = DAG.getNode(
+        LoongArchISD::BSTRPICK, DL, GRLenVT, Idx,
+        DAG.getConstant(31, DL, GRLenVT),
+        DAG.getConstant(((VecTy == MVT::v32i8) ? 2 : 1), DL, GRLenVT));
+    SDValue SplatIdx = DAG.getSplatBuildVector(MVT::v8i32, DL, NewIdx);
+    SDValue SplatValue =
+        DAG.getNode(LoongArchISD::XVPERM, DL, MVT::v8i32, NewVec, SplatIdx);
+    SDValue SplatVec = DAG.getBitcast(VecTy, SplatValue);
+
+    // Compute the local index of the original i8/i16 element within the
+    // i32 element and then use it to broadcast the vector. Each elements
+    // of the vector will be the desired element.
+    SDValue LocalIdx = DAG.getNode(
+        ISD::AND, DL, GRLenVT, Idx,
+        DAG.getConstant(((VecTy == MVT::v32i8) ? 3 : 1), DL, GRLenVT));
----------------
zhaoqi5 wrote:

Just have the vector that each i32 lane contains the desired element is not enough. We do not know which element in the i32 lane is we really want. So it is necessary to get the i32 intra-element offset(`LocalIdx`) of the desired element and use it to broadcast the element.

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