[llvm] [LoongArch] Add patterns for vstelm instructions (PR #139201)

[via llvm-commits]

https://github.com/tangaac created https://github.com/llvm/llvm-project/pull/139201

None

>From 95f8fdd66bc45fc6642fa23a7067f19d40660453 Mon Sep 17 00:00:00 2001
From: tangaac <tangyan01 at loongson.cn>
Date: Fri, 9 May 2025 11:35:53 +0800
Subject: [PATCH] add patterns for vector_extract & store merged into vstelm
 instructions

---
 .../LoongArch/LoongArchLASXInstrInfo.td       |  8 +++++++
 .../Target/LoongArch/LoongArchLSXInstrInfo.td | 21 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/llvm/lib/Target/LoongArch/LoongArchLASXInstrInfo.td b/llvm/lib/Target/LoongArch/LoongArchLASXInstrInfo.td
index fe08c1050b4d7..802fd082564e1 100644
--- a/llvm/lib/Target/LoongArch/LoongArchLASXInstrInfo.td
+++ b/llvm/lib/Target/LoongArch/LoongArchLASXInstrInfo.td
@@ -1756,6 +1756,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>;
diff --git a/llvm/lib/Target/LoongArch/LoongArchLSXInstrInfo.td b/llvm/lib/Target/LoongArch/LoongArchLSXInstrInfo.td
index 1ffc5f8056b96..69fbf5ae45603 100644
--- a/llvm/lib/Target/LoongArch/LoongArchLSXInstrInfo.td
+++ b/llvm/lib/Target/LoongArch/LoongArchLSXInstrInfo.td
@@ -1446,6 +1446,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}
@@ -1935,6 +1949,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>;