[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:

Yes, you are right. It was my misunderstanding.
I will remove it. Thanks.

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