[AArch64] Generate smov in place of sext(fmov(...))

[NickGuy]

Summary

A single smov instruction is capable of moving from a vector register while performing the sign-extend during said move, rather than each step being performed by separate instructions.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

2509–2512

These should be using VectorIndexB.

2514–2517

These should only use VectorIndexH.

2519–2520

This should be using VectorIndexS.

llvm/test/CodeGen/AArch64/aarch64-smov-gen.ll

6–18

Please simplify the tests. For example

target triple = "aarch64-unknown-linux-gnu"

define i32 @extract_s8(<vscale x 16 x i8> %a) #0 {
  %elt = extractelement <vscale x 16 x i8> %a, i32 15
  %conv = sext i8 %elt to i32
  ret i32 %conv
}

attributes #0 = { "target-features"="+sve" }

Should be enough to test the new patterns.

Given the VectorIndex# issues above I think it's worth having tests for out-of-range indices as well. I guess testing extract element VF-1 and extract element VF will cover the good and less good cases.

[NickGuy]

Addressed comments, adding additional test cases that cover cases with out-of-range indices.

[paulwalker-arm]

The out-of-range tests show that with a few more patterns we can do better, but given they're not the common case I guess they can wait.

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

2509–2520

Can you move these patterns up a couple of blocks to be just after the UMOV variants as that's the block they relate to.

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVEorStreamingSVE] in {
   def : Pat<(i32 (vector_extract (nxv8i16 ZPR:$vec), VectorIndexH:$index)),
             (i32 (UMOVvi16 (v8i16 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexH:$index))>;
   def : Pat<(i32 (vector_extract (nxv4i32 ZPR:$vec), VectorIndexS:$index)),
             (i32 (UMOVvi32 (v4i32 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexS:$index))>;
   def : Pat<(i64 (vector_extract (nxv2i64 ZPR:$vec), VectorIndexD:$index)),
             (i64 (UMOVvi64 (v2i64 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexD:$index))>;
   }
 
+  def : Pat<(sext_inreg (vector_extract (nxv16i8 ZPR:$vec), VectorIndexB:$index), i8),
+            (i32 (SMOVvi8to32 (v16i8 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexB:$index))>;
+  def : Pat<(sext_inreg (anyext (vector_extract (nxv16i8 ZPR:$vec), VectorIndexB:$index)), i8),
+            (i64 (SMOVvi8to64 (v16i8 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexB:$index))>;
+
+  def : Pat<(sext_inreg (vector_extract (nxv8i16 ZPR:$vec), VectorIndexH:$index), i16),
+            (i32 (SMOVvi16to32 (v8i16 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexH:$index))>;
+  def : Pat<(sext_inreg (anyext (vector_extract (nxv8i16 ZPR:$vec), VectorIndexH:$index)), i16),
+            (i64 (SMOVvi16to64 (v8i16 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexH:$index))>;
+
+  def : Pat<(sext (vector_extract (nxv4i32 ZPR:$vec), VectorIndexS:$index)),
+            (i64 (SMOVvi32to64 (v4i32 (EXTRACT_SUBREG ZPR:$vec, zsub)), VectorIndexS:$index))>;
+
   // Extract first element from vector.
   let AddedComplexity = 2 in {
   def : Pat<(vector_extract (nxv16i8 ZPR:$Zs), (i64 0)),
             (i32 (EXTRACT_SUBREG ZPR:$Zs, ssub))>;
   def : Pat<(vector_extract (nxv8i16 ZPR:$Zs), (i64 0)),
             (i32 (EXTRACT_SUBREG ZPR:$Zs, ssub))>;
   def : Pat<(vector_extract (nxv4i32 ZPR:$Zs), (i64 0)),
             (i32 (EXTRACT_SUBREG ZPR:$Zs, ssub))>;
let Predicates = [HasSVEorStreamingSVE] in {
   def : Pat<(nxv16i8 (vector_splice (nxv16i8 ZPR:$Z1), (nxv16i8 ZPR:$Z2), (i64 (sve_ext_imm_0_15 i32:$index)))),
             (EXT_ZZI  ZPR:$Z1, ZPR:$Z2, sve_ext_imm_0_15:$index)>;
   def : Pat<(nxv8i16 (vector_splice (nxv8i16 ZPR:$Z1), (nxv8i16 ZPR:$Z2), (i64 (sve_ext_imm_0_7 i32:$index)))),
             (EXT_ZZI  ZPR:$Z1, ZPR:$Z2, sve_ext_imm_0_7:$index)>;
   def : Pat<(nxv4i32 (vector_splice (nxv4i32 ZPR:$Z1), (nxv4i32 ZPR:$Z2), (i64 (sve_ext_imm_0_3 i32:$index)))),
             (EXT_ZZI  ZPR:$Z1, ZPR:$Z2, sve_ext_imm_0_3:$index)>;
   def : Pat<(nxv2i64 (vector_splice (nxv2i64 ZPR:$Z1), (nxv2i64 ZPR:$Z2), (i64 (sve_ext_imm_0_1 i32:$index)))),
             (EXT_ZZI  ZPR:$Z1, ZPR:$Z2, sve_ext_imm_0_1:$index)>;
+
 } // End HasSVEorStreamingSVE
 
 let Predicates = [HasSVE, HasMatMulInt8] in {
   defm  SMMLA_ZZZ : sve_int_matmul<0b00, "smmla", int_aarch64_sve_smmla>;
   defm  UMMLA_ZZZ : sve_int_matmul<0b11, "ummla", int_aarch64_sve_ummla>;
   defm USMMLA_ZZZ : sve_int_matmul<0b10, "usmmla", int_aarch64_sve_usmmla>;
 } // End HasSVE, HasMatMulInt8
 
 let Predicates = [HasSVEorStreamingSVE, HasMatMulInt8] in {
   defm USDOT_ZZZ  : sve_int_dot_mixed<"usdot", int_aarch64_sve_usdot>;
   defm USDOT_ZZZI : sve_int_dot_mixed_indexed<0, "usdot", int_aarch64_sve_usdot_lane>;
   defm SUDOT_ZZZI : sve_int_dot_mixed_indexed<1, "sudot", int_aarch64_sve_sudot_lane>;
 } // End HasSVEorStreamingSVE, HasMatMulInt8
 
 let Predicates = [HasSVE, HasMatMulFP32] in {
   defm FMMLA_ZZZ_S : sve_fp_matrix_mla<0, "fmmla", ZPR32, int_aarch64_sve_fmmla, nxv4f32>;
 } // End HasSVE, HasMatMulFP32
 
 let Predicates = [HasSVE, HasMatMulFP64] in {
   defm FMMLA_ZZZ_D : sve_fp_matrix_mla<1, "fmmla", ZPR64, int_aarch64_sve_fmmla, nxv2f64>;

llvm/test/CodeGen/AArch64/aarch64-smov-gen.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+define i32 @i8_i32(<vscale x 16 x i8> %a) #0 {
+; CHECK-LABEL: i8_i32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    smov w0, v0.b[15]
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 16 x i8> %a, i32 15
+  %conv = sext i8 %elt to i32
+  ret i32 %conv
+}
+
+define i64 @i8_i64(<vscale x 16 x i8> %a) #0 {
+; CHECK-LABEL: i8_i64:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    smov x0, v0.b[15]
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 16 x i8> %a, i32 15
+  %conv = sext i8 %elt to i64
+  ret i64 %conv
+}
+
+define i32 @i16_i32(<vscale x 8 x i16> %a) #0 {
+; CHECK-LABEL: i16_i32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    smov w0, v0.h[7]
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 8 x i16> %a, i32 7
+  %conv = sext i16 %elt to i32
+  ret i32 %conv
+}
+
+define i64 @i16_i64(<vscale x 8 x i16> %a) #0 {
+; CHECK-LABEL: i16_i64:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    smov x0, v0.h[7]
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 8 x i16> %a, i32 7
+  %conv = sext i16 %elt to i64
+  ret i64 %conv
+}
+
+define i64 @i32_i64(<vscale x 4 x i32> %a) #0 {
+; CHECK-LABEL: i32_i64:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    smov x0, v0.s[3]
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 4 x i32> %a, i32 3
+  %conv = sext i32 %elt to i64
+  ret i64 %conv
+}
+
+; NOTE: Testing out-of-range indices
+
+define i32 @i8_i32_oor(<vscale x 16 x i8> %a) #0 {
+; CHECK-LABEL: i8_i32_oor:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov z0.b, z0.b[16]
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    sxtb w0, w8
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 16 x i8> %a, i32 16
+  %conv = sext i8 %elt to i32
+  ret i32 %conv
+}
+
+define i64 @i8_i64_oor(<vscale x 16 x i8> %a) #0 {
+; CHECK-LABEL: i8_i64_oor:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov z0.b, z0.b[16]
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    sxtb x0, w8
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 16 x i8> %a, i32 16
+  %conv = sext i8 %elt to i64
+  ret i64 %conv
+}
+
+define i32 @i16_i32_oor(<vscale x 8 x i16> %a) #0 {
+; CHECK-LABEL: i16_i32_oor:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov z0.h, z0.h[8]
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    sxth w0, w8
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 8 x i16> %a, i32 8
+  %conv = sext i16 %elt to i32
+  ret i32 %conv
+}
+
+define i64 @i16_i64_oor(<vscale x 8 x i16> %a) #0 {
+; CHECK-LABEL: i16_i64_oor:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov z0.h, z0.h[8]
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    sxth x0, w8
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 8 x i16> %a, i32 8
+  %conv = sext i16 %elt to i64
+  ret i64 %conv
+}
+
+define i64 @i32_i64_oor(<vscale x 4 x i32> %a) #0 {
+; CHECK-LABEL: i32_i64_oor:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov z0.s, z0.s[4]
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    sxtw x0, w8
+; CHECK-NEXT:    ret
+entry:
+  %elt = extractelement <vscale x 4 x i32> %a, i32 4
+  %conv = sext i32 %elt to i64
+  ret i64 %conv
+}
+
+attributes #0 = { "target-features"="+sve" }