add patterns for vector_extract & store merged into vstelm instructions

Author: tangaac

llvm/lib/Target/LoongArch/LoongArchLASXInstrInfo.td

@@ -1760,6 +1760,14 @@ def : Pat<(lasxsplatf32 FPR32:$fj),
 def : Pat<(lasxsplatf64 FPR64:$fj),
           (XVREPLVE0_D (SUBREG_TO_REG (i64 0), FPR64:$fj, sub_64))>;
 
+// VSTELM
+defm : VstelmPat<truncstorei8, v32i8, XVSTELM_B, simm12_addlike, uimm5>;
+defm : VstelmPat<truncstorei16, v16i16, XVSTELM_H, simm11_lsl1, uimm4>;
+defm : VstelmPat<truncstorei32, v8i32, XVSTELM_W, simm10_lsl2, uimm3>;
+defm : VstelmPat<store, v4i64, XVSTELM_D, simm9_lsl3, uimm2>;
+defm : VstelmPat<store, v8f32, XVSTELM_W, simm10_lsl2, uimm3, f32>;
+defm : VstelmPat<store, v4f64, XVSTELM_D, simm9_lsl3, uimm2, f64>;
+
 // Loads/Stores
 foreach vt = [v32i8, v16i16, v8i32, v4i64, v8f32, v4f64] in {
   defm : LdPat<load, XVLD, vt>;

llvm/lib/Target/LoongArch/LoongArchLSXInstrInfo.td

@@ -1452,6 +1452,20 @@ multiclass VldreplPat<ValueType vt, LAInst Inst, Operand ImmOpnd> {
             (Inst BaseAddr:$rj, ImmOpnd:$imm)>;
 }
 
+multiclass VstelmPat<PatFrag StoreOp, ValueType vt, LAInst Inst,
+                     Operand ImmOpnd, Operand IdxOpnd, ValueType elt = i64> {
+  def : Pat<(StoreOp(elt(vector_extract vt:$vd, IdxOpnd:$idx)), BaseAddr:$rj),
+            (Inst vt:$vd, BaseAddr:$rj, 0, IdxOpnd:$idx)>;
+
+  def : Pat<(StoreOp(elt(vector_extract vt:$vd, IdxOpnd:$idx)),
+                (AddrConstant GPR:$rj, ImmOpnd:$imm)),
+            (Inst vt:$vd, GPR:$rj, ImmOpnd:$imm, IdxOpnd:$idx)>;
+
+  def : Pat<(StoreOp(elt(vector_extract vt:$vd, IdxOpnd:$idx)),
+                (AddLike BaseAddr:$rj, ImmOpnd:$imm)),
+            (Inst vt:$vd, BaseAddr:$rj, ImmOpnd:$imm, IdxOpnd:$idx)>;
+}
+
 let Predicates = [HasExtLSX] in {
 
 // VADD_{B/H/W/D}
@@ -1945,6 +1959,13 @@ def : Pat<(lsxsplatf32 FPR32:$fj),
 def : Pat<(lsxsplatf64 FPR64:$fj),
           (VREPLVEI_D (SUBREG_TO_REG (i64 0), FPR64:$fj, sub_64), 0)>;
 
+defm : VstelmPat<truncstorei8, v16i8, VSTELM_B, simm12_addlike, uimm4>;
+defm : VstelmPat<truncstorei16, v8i16, VSTELM_H, simm11_lsl1, uimm3>;
+defm : VstelmPat<truncstorei32, v4i32, VSTELM_W, simm10_lsl2, uimm2>;
+defm : VstelmPat<store, v2i64, VSTELM_D, simm9_lsl3, uimm1>;
+defm : VstelmPat<store, v4f32, VSTELM_W, simm10_lsl2, uimm2, f32>;
+defm : VstelmPat<store, v2f64, VSTELM_D, simm9_lsl3, uimm1, f64>;
+
 // Loads/Stores
 foreach vt = [v16i8, v8i16, v4i32, v2i64, v4f32, v2f64] in {
   defm : LdPat<load, VLD, vt>;