[RISCV] Supplement patterns for vnsrl.wx/vnsra.wx when splat shift is sext or zext

[Jimerlife]

Summary

Supplement this patch D102019, I think when splat shift is sext or zext also should match vnsrl.wx or vnsra.wx。

[craig.topper]

LGTM

llvm/lib/Target/RISCV/RISCVInstrInfoVVLPatterns.td

multiclass VPatWidenBinaryFPVL_WV_WF<SDNode op, PatFrags extop, string instruction_name> {
   }
 }
 
 multiclass VPatWidenBinaryFPVL_VV_VF_WV_WF<SDNode op, string instruction_name> {
   defm : VPatWidenBinaryFPVL_VV_VF<op, riscv_fpextend_vl_oneuse, instruction_name>;
   defm : VPatWidenBinaryFPVL_WV_WF<op, riscv_fpextend_vl_oneuse, instruction_name>;
 }
 
+multiclass VPatNarrowShiftSplatExt_WX<SDNode op, PatFrags extop, string instruction_name> {
+  foreach vti = AllWidenableIntVectors in {
+    def : Pat<
+      (vti.Vti.Vector
+        (riscv_trunc_vector_vl
+          (op (vti.Wti.Vector vti.Wti.RegClass:$rs2),
+              (vti.Wti.Vector (extop (vti.Vti.Vector (SplatPat GPR:$rs1)),
+                                     (vti.Vti.Mask true_mask), VLOpFrag)),
+          (vti.Wti.Mask true_mask), VLOpFrag),
+        (vti.Vti.Mask true_mask), VLOpFrag)),
+      (!cast<Instruction>(instruction_name#"_WX_"#vti.Vti.LMul.MX)
+        vti.Wti.RegClass:$rs2, GPR:$rs1,
+        vti.Vti.AVL, vti.Vti.Log2SEW)>;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Patterns.
 //===----------------------------------------------------------------------===//
 
 let Predicates = [HasVInstructions] in {
 
 // 12.1. Vector Single-Width Integer Add and Subtract
 defm : VPatBinaryVL_VV_VX_VI<riscv_add_vl, "PseudoVADD">;
foreach vtiTowti = AllWidenableIntVectors in {
 
   def : Pat<(vti.Vector
              (riscv_trunc_vector_vl
               (wti.Vector
                (riscv_sra_vl wti.RegClass:$rs1, (SplatPat XLenVT:$rs2),
                              true_mask, VLOpFrag)), true_mask, VLOpFrag)),
             (!cast<Instruction>("PseudoVNSRA_WX_"#vti.LMul.MX)
                 wti.RegClass:$rs1, GPR:$rs2, GPR:$vl, vti.Log2SEW)>;
+  defm :  VPatNarrowShiftSplatExt_WX<riscv_sra_vl, riscv_sext_vl_oneuse, "PseudoVNSRA">;
+  defm :  VPatNarrowShiftSplatExt_WX<riscv_sra_vl, riscv_zext_vl_oneuse, "PseudoVNSRA">;
   def : Pat<(vti.Vector
              (riscv_trunc_vector_vl
               (wti.Vector
                (riscv_sra_vl wti.RegClass:$rs1, (SplatPat_uimm5 uimm5:$rs2),
                              true_mask, VLOpFrag)), true_mask, VLOpFrag)),
             (!cast<Instruction>("PseudoVNSRA_WI_"#vti.LMul.MX)
                 wti.RegClass:$rs1, uimm5:$rs2, GPR:$vl, vti.Log2SEW)>;
-
   def : Pat<(vti.Vector
              (riscv_trunc_vector_vl
               (wti.Vector
                (riscv_srl_vl wti.RegClass:$rs1, (SplatPat XLenVT:$rs2),
                              true_mask, VLOpFrag)), true_mask, VLOpFrag)),
             (!cast<Instruction>("PseudoVNSRL_WX_"#vti.LMul.MX)
                 wti.RegClass:$rs1, GPR:$rs2, GPR:$vl, vti.Log2SEW)>;
+  defm :  VPatNarrowShiftSplatExt_WX<riscv_srl_vl, riscv_sext_vl_oneuse, "PseudoVNSRL">;
+  defm :  VPatNarrowShiftSplatExt_WX<riscv_srl_vl, riscv_zext_vl_oneuse, "PseudoVNSRL">;
   def : Pat<(vti.Vector
              (riscv_trunc_vector_vl
               (wti.Vector
                (riscv_srl_vl wti.RegClass:$rs1, (SplatPat_uimm5 uimm5:$rs2),
                              true_mask, VLOpFrag)), true_mask, VLOpFrag)),
             (!cast<Instruction>("PseudoVNSRL_WI_"#vti.LMul.MX)
                 wti.RegClass:$rs1, uimm5:$rs2, GPR:$vl, vti.Log2SEW)>;
 }

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vnsra-vnsrl.ll

 ; CHECK-NEXT:    ret
   %insert = insertelement <8 x i16> poison, i16 %y, i16 0
   %splat = shufflevector <8 x i16> %insert, <8 x i16> poison, <8 x i32> zeroinitializer
   %a = ashr <8 x i16> %x, %splat
   %b = trunc <8 x i16> %a to <8 x i8>
   ret <8 x i8> %b
 }
 
+define <8 x i8> @vnsra_v8i16_v8i8_scalar_sext(<8 x i16> %x, i8 %y) {
+; CHECK-LABEL: vnsra_v8i16_v8i8_scalar_sext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e8, mf2, ta, mu
+; CHECK-NEXT:    vnsra.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <8 x i8> poison, i8 %y, i8 0
+  %splat = shufflevector <8 x i8> %insert, <8 x i8> poison, <8 x i32> zeroinitializer
+  %sext = sext <8 x i8> %splat to <8 x i16>
+  %a = ashr <8 x i16> %x, %sext
+  %b = trunc <8 x i16> %a to <8 x i8>
+  ret <8 x i8> %b
+}
+
+define <8 x i8> @vnsra_v8i16_v8i8_scalar_zext(<8 x i16> %x, i8 %y) {
+; CHECK-LABEL: vnsra_v8i16_v8i8_scalar_zext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e8, mf2, ta, mu
+; CHECK-NEXT:    vnsra.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <8 x i8> poison, i8 %y, i8 0
+  %splat = shufflevector <8 x i8> %insert, <8 x i8> poison, <8 x i32> zeroinitializer
+  %zext = zext <8 x i8> %splat to <8 x i16>
+  %a = ashr <8 x i16> %x, %zext
+  %b = trunc <8 x i16> %a to <8 x i8>
+  ret <8 x i8> %b
+}
+
 define <4 x i16> @vnsra_v4i32_v4i16_scalar(<4 x i32> %x, i32 %y) {
 ; CHECK-LABEL: vnsra_v4i32_v4i16_scalar:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 4, e16, mf2, ta, mu
 ; CHECK-NEXT:    vnsra.wx v8, v8, a0
 ; CHECK-NEXT:    ret
   %insert = insertelement <4 x i32> poison, i32 %y, i32 0
   %splat = shufflevector <4 x i32> %insert, <4 x i32> poison, <4 x i32> zeroinitializer
   %a = ashr <4 x i32> %x, %splat
   %b = trunc <4 x i32> %a to <4 x i16>
   ret <4 x i16> %b
 }
 
+define <4 x i16> @vnsra_v4i32_v4i16_scalar_sext(<4 x i32> %x, i16 %y) {
+; CHECK-LABEL: vnsra_v4i32_v4i16_scalar_sext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
+; CHECK-NEXT:    vnsra.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <4 x i16> poison, i16 %y, i16 0
+  %splat = shufflevector <4 x i16> %insert, <4 x i16> poison, <4 x i32> zeroinitializer
+  %sext = sext <4 x i16> %splat to <4 x i32>
+  %a = ashr <4 x i32> %x, %sext
+  %b = trunc <4 x i32> %a to <4 x i16>
+  ret <4 x i16> %b
+}
+
+define <4 x i16> @vnsra_v4i32_v4i16_scalar_zext(<4 x i32> %x, i16 %y) {
+; CHECK-LABEL: vnsra_v4i32_v4i16_scalar_zext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
+; CHECK-NEXT:    vnsra.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <4 x i16> poison, i16 %y, i16 0
+  %splat = shufflevector <4 x i16> %insert, <4 x i16> poison, <4 x i32> zeroinitializer
+  %zext = zext <4 x i16> %splat to <4 x i32>
+  %a = ashr <4 x i32> %x, %zext
+  %b = trunc <4 x i32> %a to <4 x i16>
+  ret <4 x i16> %b
+}
+
 define <2 x i32> @vnsra_v2i64_v2i32_scalar(<2 x i64> %x, i64 %y) {
 ; CHECK-LABEL: vnsra_v2i64_v2i32_scalar:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 2, e32, mf2, ta, mu
 ; CHECK-NEXT:    vnsra.wx v8, v8, a0
 ; CHECK-NEXT:    ret
   %insert = insertelement <2 x i64> poison, i64 %y, i32 0
   %splat = shufflevector <2 x i64> %insert, <2 x i64> poison, <2 x i32> zeroinitializer
   %a = ashr <2 x i64> %x, %splat
   %b = trunc <2 x i64> %a to <2 x i32>
   ret <2 x i32> %b
 }
 
+define <2 x i32> @vnsra_v2i64_v2i32_scalar_sext(<2 x i64> %x, i32 %y) {
+; CHECK-LABEL: vnsra_v2i64_v2i32_scalar_sext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e32, mf2, ta, mu
+; CHECK-NEXT:    vnsra.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <2 x i32> poison, i32 %y, i32 0
+  %splat = shufflevector <2 x i32> %insert, <2 x i32> poison, <2 x i32> zeroinitializer
+  %sext = sext <2 x i32> %splat to <2 x i64>
+  %a = ashr <2 x i64> %x, %sext
+  %b = trunc <2 x i64> %a to <2 x i32>
+  ret <2 x i32> %b
+}
+
+define <2 x i32> @vnsra_v2i64_v2i32_scalar_zext(<2 x i64> %x, i32 %y) {
+; CHECK-LABEL: vnsra_v2i64_v2i32_scalar_zext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e32, mf2, ta, mu
+; CHECK-NEXT:    vnsra.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <2 x i32> poison, i32 %y, i32 0
+  %splat = shufflevector <2 x i32> %insert, <2 x i32> poison, <2 x i32> zeroinitializer
+  %zext = zext <2 x i32> %splat to <2 x i64>
+  %a = ashr <2 x i64> %x, %zext
+  %b = trunc <2 x i64> %a to <2 x i32>
+  ret <2 x i32> %b
+}
+
 define <8 x i8> @vnsra_v8i16_v8i8_imm(<8 x i16> %x) {
 ; CHECK-LABEL: vnsra_v8i16_v8i8_imm:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, mu
 ; CHECK-NEXT:    vnsrl.wi v8, v8, 8
 ; CHECK-NEXT:    ret
   %a = ashr <8 x i16> %x, <i16 8, i16 8, i16 8, i16 8,i16 8, i16 8, i16 8, i16 8>
   %b = trunc <8 x i16> %a to <8 x i8>

 ; CHECK-NEXT:    ret
   %insert = insertelement <8 x i16> poison, i16 %y, i16 0
   %splat = shufflevector <8 x i16> %insert, <8 x i16> poison, <8 x i32> zeroinitializer
   %a = lshr <8 x i16> %x, %splat
   %b = trunc <8 x i16> %a to <8 x i8>
   ret <8 x i8> %b
 }
 
+define <8 x i8> @vnsrl_v8i16_v8i8_scalar_sext(<8 x i16> %x, i8 %y) {
+; CHECK-LABEL: vnsrl_v8i16_v8i8_scalar_sext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e8, mf2, ta, mu
+; CHECK-NEXT:    vnsrl.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <8 x i8> poison, i8 %y, i16 0
+  %splat = shufflevector <8 x i8> %insert, <8 x i8> poison, <8 x i32> zeroinitializer
+  %sext = sext <8 x i8> %splat to <8 x i16>
+  %a = lshr <8 x i16> %x, %sext
+  %b = trunc <8 x i16> %a to <8 x i8>
+  ret <8 x i8> %b
+}
+
+define <8 x i8> @vnsrl_v8i16_v8i8_scalar_zext(<8 x i16> %x, i8 %y) {
+; CHECK-LABEL: vnsrl_v8i16_v8i8_scalar_zext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e8, mf2, ta, mu
+; CHECK-NEXT:    vnsrl.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <8 x i8> poison, i8 %y, i16 0
+  %splat = shufflevector <8 x i8> %insert, <8 x i8> poison, <8 x i32> zeroinitializer
+  %zext = zext <8 x i8> %splat to <8 x i16>
+  %a = lshr <8 x i16> %x, %zext
+  %b = trunc <8 x i16> %a to <8 x i8>
+  ret <8 x i8> %b
+}
+
 define <4 x i16> @vnsrl_v4i32_v4i16_scalar(<4 x i32> %x, i32 %y) {
 ; CHECK-LABEL: vnsrl_v4i32_v4i16_scalar:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 4, e16, mf2, ta, mu
 ; CHECK-NEXT:    vnsrl.wx v8, v8, a0
 ; CHECK-NEXT:    ret
   %insert = insertelement <4 x i32> poison, i32 %y, i32 0
   %splat = shufflevector <4 x i32> %insert, <4 x i32> poison, <4 x i32> zeroinitializer
   %a = lshr <4 x i32> %x, %splat
   %b = trunc <4 x i32> %a to <4 x i16>
   ret <4 x i16> %b
 }
 
+define <4 x i16> @vnsrl_v4i32_v4i16_scalar_sext(<4 x i32> %x, i16 %y) {
+; CHECK-LABEL: vnsrl_v4i32_v4i16_scalar_sext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
+; CHECK-NEXT:    vnsrl.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <4 x i16> poison, i16 %y, i16 0
+  %splat = shufflevector <4 x i16> %insert, <4 x i16> poison, <4 x i32> zeroinitializer
+  %sext = sext <4 x i16> %splat to <4 x i32>
+  %a = lshr <4 x i32> %x, %sext
+  %b = trunc <4 x i32> %a to <4 x i16>
+  ret <4 x i16> %b
+}
+
+define <4 x i16> @vnsrl_v4i32_v4i16_scalar_zext(<4 x i32> %x, i16 %y) {
+; CHECK-LABEL: vnsrl_v4i32_v4i16_scalar_zext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
+; CHECK-NEXT:    vnsrl.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <4 x i16> poison, i16 %y, i16 0
+  %splat = shufflevector <4 x i16> %insert, <4 x i16> poison, <4 x i32> zeroinitializer
+  %zext = zext <4 x i16> %splat to <4 x i32>
+  %a = lshr <4 x i32> %x, %zext
+  %b = trunc <4 x i32> %a to <4 x i16>
+  ret <4 x i16> %b
+}
+
 define <2 x i32> @vnsrl_v2i64_v2i32_scalar(<2 x i64> %x, i64 %y) {
 ; CHECK-LABEL: vnsrl_v2i64_v2i32_scalar:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 2, e32, mf2, ta, mu
 ; CHECK-NEXT:    vnsrl.wx v8, v8, a0
 ; CHECK-NEXT:    ret
   %insert = insertelement <2 x i64> poison, i64 %y, i32 0
   %splat = shufflevector <2 x i64> %insert, <2 x i64> poison, <2 x i32> zeroinitializer
   %a = lshr <2 x i64> %x, %splat
   %b = trunc <2 x i64> %a to <2 x i32>
   ret <2 x i32> %b
 }
 
+define <2 x i32> @vnsrl_v2i64_v2i32_scalar_sext(<2 x i64> %x, i32 %y) {
+; CHECK-LABEL: vnsrl_v2i64_v2i32_scalar_sext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e32, mf2, ta, mu
+; CHECK-NEXT:    vnsrl.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <2 x i32> poison, i32 %y, i32 0
+  %splat = shufflevector <2 x i32> %insert, <2 x i32> poison, <2 x i32> zeroinitializer
+  %sext = sext <2 x i32> %splat to <2 x i64>
+  %a = lshr <2 x i64> %x, %sext
+  %b = trunc <2 x i64> %a to <2 x i32>
+  ret <2 x i32> %b
+}
+
+define <2 x i32> @vnsrl_v2i64_v2i32_scalar_zext(<2 x i64> %x, i32 %y) {
+; CHECK-LABEL: vnsrl_v2i64_v2i32_scalar_zext:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e32, mf2, ta, mu
+; CHECK-NEXT:    vnsrl.wx v8, v8, a0
+; CHECK-NEXT:    ret
+  %insert = insertelement <2 x i32> poison, i32 %y, i32 0
+  %splat = shufflevector <2 x i32> %insert, <2 x i32> poison, <2 x i32> zeroinitializer
+  %zext = zext <2 x i32> %splat to <2 x i64>
+  %a = lshr <2 x i64> %x, %zext
+  %b = trunc <2 x i64> %a to <2 x i32>
+  ret <2 x i32> %b
+}
+
 define <8 x i8> @vnsrl_v8i16_v8i8_imm(<8 x i16> %x) {
 ; CHECK-LABEL: vnsrl_v8i16_v8i8_imm:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, mu
 ; CHECK-NEXT:    vnsrl.wi v8, v8, 8
 ; CHECK-NEXT:    ret
   %a = lshr <8 x i16> %x, <i16 8, i16 8, i16 8, i16 8,i16 8, i16 8, i16 8, i16 8>
   %b = trunc <8 x i16> %a to <8 x i8>