[AArch64] Improve single vector lane stores

[evandro]

Summary

When storing the 0th lane of a vector, use a simpler and usually more efficient scalar store instead.

[sebpop]

This fixes a perf regression we were seeing with generation of vectors smaller than 64 bit: https://reviews.llvm.org/D45821

[efriedma]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2385

Do we also need patterns for STURB and friends?

[evandro]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2385

I came across some code indicating that we do. But I'm working on that in a different patch, testing pending.

[efriedma]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2385

Okay; please add some test coverage for the offsets we current do/don't handle.

[efriedma]

llvm/test/CodeGen/AArch64/arm64-st1.ll

51

Please don't use CHECK-NOT like this; it isn't clear what we're actually generating.

[efriedma]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2256

The VecROStoreLane0Pat<ro8, store, v16i8, VecROStoreLane0Pat<ro16, store, v8i16, and VecROStoreLane0Pat<ro32, truncstorei32, v4i32 patterns will never match. And there's a missing truncstorei8 pattern.

[evandro]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2256

True, except for VecROStoreLane0Pat<ro8, ..., since FPR8, used by STRBro{W,X}, can only contain untyped.

[evandro]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2256

Or do you mean like so?

` defm : VecStoreULane0Pat<truncstorei8, v8i16, i32, hsub, STURBi>;

defm : VecStoreULane0Pat<truncstorei16, v8i16, i32, hsub, STURHi>;
defm : VecStoreULane0Pat<store,         v8f16, f16, hsub, STURHi>;
defm : VecStoreULane0Pat<truncstorei8,  v4i32, i32, ssub, STURBi>;
defm : VecStoreULane0Pat<truncstorei16, v4i32, i32, ssub, STURHi>;
defm : VecStoreULane0Pat<truncstorei32, v4i32, i32, ssub, STURSi>;`

[efriedma]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2256

I'm not sure where you're going with your suggestion.

Due to legalization, the result of a vector_extract can't be i8 or i16; the patterns need to account for that.

[evandro]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2256

I'll update this patch and then, please, let me know if I'm following what you mean.

[efriedma]

LGTM. (We should also add the i8 patterns at some point, but we can do that as a followup.)

[evandro]

In D46655#1096661, @efriedma wrote:

LGTM. (We should also add the i8 patterns at some point, but we can do that as a followup.)

Trying to add i8 hasn't been as straightforward. It seems to be hinged on FPR8 holding only untyped, but the approach that I'm trying would not change that. I'll keep you posted.

Thank you.

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 }
 
 let AddedComplexity = 19 in {
+  defm : VecROStoreLane0Pat<ro8,          store, v16i8,  i8, bsub, STRBroW, STRBroX>;
   defm : VecROStoreLane0Pat<ro16, truncstorei16, v8i16, i32, hsub, STRHroW, STRHroX>;
-  defm : VecROStoreLane0Pat<ro16,      store   , v8i16, i16, hsub, STRHroW, STRHroX>;
+  defm : VecROStoreLane0Pat<ro16,         store, v8i16, i16, hsub, STRHroW, STRHroX>;
+  defm : VecROStoreLane0Pat<ro16,         store, v8f16, f16, hsub, STRHroW, STRHroX>;
   defm : VecROStoreLane0Pat<ro32, truncstorei32, v4i32, i32, ssub, STRSroW, STRSroX>;
-  defm : VecROStoreLane0Pat<ro32,      store   , v4i32, i32, ssub, STRSroW, STRSroX>;
-  defm : VecROStoreLane0Pat<ro32,      store   , v4f32, f32, ssub, STRSroW, STRSroX>;
-  defm : VecROStoreLane0Pat<ro64,      store   , v2i64, i64, dsub, STRDroW, STRDroX>;
-  defm : VecROStoreLane0Pat<ro64,      store   , v2f64, f64, dsub, STRDroW, STRDroX>;
+  defm : VecROStoreLane0Pat<ro32,         store, v4i32, i32, ssub, STRSroW, STRSroX>;
+  defm : VecROStoreLane0Pat<ro32,         store, v4f32, f32, ssub, STRSroW, STRSroX>;
+  defm : VecROStoreLane0Pat<ro64,         store, v2i64, i64, dsub, STRDroW, STRDroX>;
+  defm : VecROStoreLane0Pat<ro64,         store, v2f64, f64, dsub, STRDroW, STRDroX>;
 }
 
 //---

                                     (am_indexed8 GPR64sp:$Rn,
                                                  uimm12s1:$offset))]>;
 
-// Match all store 64 bits width whose type is compatible with FPR64
 let AddedComplexity = 10 in {
+
+// Match all store 64 bits width whose type is compatible with FPR64
+def : Pat<(store (v1i64 FPR64:$Rt),
+                 (am_indexed64 GPR64sp:$Rn, uimm12s8:$offset)),
+          (STRDui FPR64:$Rt, GPR64sp:$Rn, uimm12s8:$offset)>;
+def : Pat<(store (v1f64 FPR64:$Rt),
+                 (am_indexed64 GPR64sp:$Rn, uimm12s8:$offset)),
+          (STRDui FPR64:$Rt, GPR64sp:$Rn, uimm12s8:$offset)>;
+
 let Predicates = [IsLE] in {
   // We must use ST1 to store vectors in big-endian.
   def : Pat<(store (v2f32 FPR64:$Rt),

                    (am_indexed64 GPR64sp:$Rn, uimm12s8:$offset)),
             (STRDui FPR64:$Rt, GPR64sp:$Rn, uimm12s8:$offset)>;
 }
-def : Pat<(store (v1f64 FPR64:$Rt),
-                 (am_indexed64 GPR64sp:$Rn, uimm12s8:$offset)),
-          (STRDui FPR64:$Rt, GPR64sp:$Rn, uimm12s8:$offset)>;
-def : Pat<(store (v1i64 FPR64:$Rt),
-                 (am_indexed64 GPR64sp:$Rn, uimm12s8:$offset)),
-          (STRDui FPR64:$Rt, GPR64sp:$Rn, uimm12s8:$offset)>;
 
 // Match all store 128 bits width whose type is compatible with FPR128
+def : Pat<(store (f128  FPR128:$Rt),
+                 (am_indexed128 GPR64sp:$Rn, uimm12s16:$offset)),
+          (STRQui FPR128:$Rt, GPR64sp:$Rn, uimm12s16:$offset)>;
+
 let Predicates = [IsLE] in {
   // We must use ST1 to store vectors in big-endian.
   def : Pat<(store (v4f32 FPR128:$Rt),

                    (am_indexed128 GPR64sp:$Rn, uimm12s16:$offset)),
             (STRQui FPR128:$Rt, GPR64sp:$Rn, uimm12s16:$offset)>;
 }
-def : Pat<(store (f128  FPR128:$Rt),
-                 (am_indexed128 GPR64sp:$Rn, uimm12s16:$offset)),
-          (STRQui FPR128:$Rt, GPR64sp:$Rn, uimm12s16:$offset)>;
 
 // truncstore i64
 def : Pat<(truncstorei32 GPR64:$Rt,

 
 } // AddedComplexity = 10
 
+// Match stores from lane 0 to the appropriate subreg's store.
+multiclass VecStoreLane0Pat<Operand UIAddrMode, SDPatternOperator storeop,
+                            ValueType VTy, ValueType STy,
+                            SubRegIndex SubRegIdx, Operand IndexType,
+                            Instruction STR> {
+  def : Pat<(storeop (STy (vector_extract (VTy VecListOne128:$Vt), 0)),
+                     (UIAddrMode GPR64sp:$Rn, IndexType:$offset)),
+            (STR (EXTRACT_SUBREG VecListOne128:$Vt, SubRegIdx),
+                 GPR64sp:$Rn, IndexType:$offset)>;
+}
+
+let AddedComplexity = 19 in {
+  defm : VecStoreLane0Pat<am_indexed8,          store, v16i8,  i8, bsub, uimm12s1, STRBui>;
+  defm : VecStoreLane0Pat<am_indexed16, truncstorei16, v8i16, i32, hsub, uimm12s2, STRHui>;
+  defm : VecStoreLane0Pat<am_indexed16,         store, v8i16, i16, hsub, uimm12s2, STRHui>;
+  defm : VecStoreLane0Pat<am_indexed16,         store, v8f16, f16, hsub, uimm12s2, STRHui>;
+  defm : VecStoreLane0Pat<am_indexed32, truncstorei32, v4i32, i32, ssub, uimm12s4, STRSui>;
+  defm : VecStoreLane0Pat<am_indexed32,         store, v4i32, i32, ssub, uimm12s4, STRSui>;
+  defm : VecStoreLane0Pat<am_indexed32,         store, v4f32, f32, ssub, uimm12s4, STRSui>;
+  defm : VecStoreLane0Pat<am_indexed64,         store, v2i64, i64, dsub, uimm12s8, STRDui>;
+  defm : VecStoreLane0Pat<am_indexed64,         store, v2f64, f64, dsub, uimm12s8, STRDui>;
+}
+
 //---
 // (unscaled immediate)
 defm STURX : StoreUnscaled<0b11, 0, 0b00, GPR64, "stur",

                                   (am_unscaled8 GPR64sp:$Rn, simm9:$offset))]>;
 
 // Match all store 64 bits width whose type is compatible with FPR64
+def : Pat<(store (v1f64 FPR64:$Rt), (am_unscaled64 GPR64sp:$Rn, simm9:$offset)),
+          (STURDi FPR64:$Rt, GPR64sp:$Rn, simm9:$offset)>;
+def : Pat<(store (v1i64 FPR64:$Rt), (am_unscaled64 GPR64sp:$Rn, simm9:$offset)),
+          (STURDi FPR64:$Rt, GPR64sp:$Rn, simm9:$offset)>;
+
 let Predicates = [IsLE] in {
   // We must use ST1 to store vectors in big-endian.
   def : Pat<(store (v2f32 FPR64:$Rt),

                    (am_unscaled64 GPR64sp:$Rn, simm9:$offset)),
             (STURDi FPR64:$Rt, GPR64sp:$Rn, simm9:$offset)>;
 }
-def : Pat<(store (v1f64 FPR64:$Rt), (am_unscaled64 GPR64sp:$Rn, simm9:$offset)),
-          (STURDi FPR64:$Rt, GPR64sp:$Rn, simm9:$offset)>;
-def : Pat<(store (v1i64 FPR64:$Rt), (am_unscaled64 GPR64sp:$Rn, simm9:$offset)),
-          (STURDi FPR64:$Rt, GPR64sp:$Rn, simm9:$offset)>;
 
 // Match all store 128 bits width whose type is compatible with FPR128
+def : Pat<(store (f128 FPR128:$Rt), (am_unscaled128 GPR64sp:$Rn, simm9:$offset)),
+          (STURQi FPR128:$Rt, GPR64sp:$Rn, simm9:$offset)>;
+
 let Predicates = [IsLE] in {
   // We must use ST1 to store vectors in big-endian.
   def : Pat<(store (v4f32 FPR128:$Rt),

           (SUBREG_TO_REG (i32 0),
                          (f32 (COPY_TO_REGCLASS GPR32:$Rn, FPR32)), ssub)>;
 
+def : Pat<(v4f16 (scalar_to_vector (f16 FPR16:$Rn))),
+          (INSERT_SUBREG (v4f16 (IMPLICIT_DEF)), FPR16:$Rn, hsub)>;
+def : Pat<(v8f16 (scalar_to_vector (f16 FPR16:$Rn))),
+          (INSERT_SUBREG (v8f16 (IMPLICIT_DEF)), FPR16:$Rn, hsub)>;
+
 def : Pat<(v2i32 (scalar_to_vector (i32 FPR32:$Rn))),
             (v2i32 (INSERT_SUBREG (v2i32 (IMPLICIT_DEF)),
                                   (i32 FPR32:$Rn), ssub))>;
 def : Pat<(v4i32 (scalar_to_vector (i32 FPR32:$Rn))),
             (v4i32 (INSERT_SUBREG (v4i32 (IMPLICIT_DEF)),
                                   (i32 FPR32:$Rn), ssub))>;
+
 def : Pat<(v2i64 (scalar_to_vector (i64 FPR64:$Rn))),
             (v2i64 (INSERT_SUBREG (v2i64 (IMPLICIT_DEF)),
                                   (i64 FPR64:$Rn), dsub))>;

           (INSERT_SUBREG (v4f32 (IMPLICIT_DEF)), FPR32:$Rn, ssub)>;
 def : Pat<(v2f32 (scalar_to_vector (f32 FPR32:$Rn))),
           (INSERT_SUBREG (v2f32 (IMPLICIT_DEF)), FPR32:$Rn, ssub)>;
+
 def : Pat<(v2f64 (scalar_to_vector (f64 FPR64:$Rn))),
           (INSERT_SUBREG (v2f64 (IMPLICIT_DEF)), FPR64:$Rn, dsub)>;
 

llvm/test/CodeGen/AArch64/arm64-neon-copy.ll

 ; CHECK-LABEL: test_extracts_inserts_varidx_insert:
 ; CHECK: and [[MASKED_IDX:x[0-9]+]], x0, #0x3
 ; CHECK: bfi x9, [[MASKED_IDX]], #1, #2
-; CHECK: st1 { v0.h }[0], [x9]
+; CHECK: str h0, [x9]
 ; CHECK-DAG: ldr d[[R:[0-9]+]]
 ; CHECK-DAG: mov v[[R]].h[1], v0.h[1]
 ; CHECK-DAG: mov v[[R]].h[2], v0.h[2]

llvm/test/CodeGen/AArch64/arm64-neon-simd-ldst-one.ll

   ret void
 }
 
+define void @test_vst1q_lane0_s16(i16* %a, <8 x i16> %b) {
+; CHECK-LABEL: test_vst1q_lane0_s16:
+; CHECK: str {{h[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <8 x i16> %b, i32 0
+  store i16 %0, i16* %a, align 2
+  ret void
+}
+
 define void @test_vst1q_lane_s32(i32* %a, <4 x i32> %b) {
 ; CHECK-LABEL: test_vst1q_lane_s32:
 ; CHECK: st1 { {{v[0-9]+}}.s }[{{[0-9]+}}], [x0]

   ret void
 }
 
+define void @test_vst1q_lane0_s32(i32* %a, <4 x i32> %b) {
+; CHECK-LABEL: test_vst1q_lane0_s32:
+; CHECK: str {{s[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <4 x i32> %b, i32 0
+  store i32 %0, i32* %a, align 4
+  ret void
+}
+
 define void @test_vst1q_lane_s64(i64* %a, <2 x i64> %b) {
 ; CHECK-LABEL: test_vst1q_lane_s64:
 ; CHECK: st1 { {{v[0-9]+}}.d }[{{[0-9]+}}], [x0]

   ret void
 }
 
+define void @test_vst1q_lane0_s64(i64* %a, <2 x i64> %b) {
+; CHECK-LABEL: test_vst1q_lane0_s64:
+; CHECK: str {{d[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <2 x i64> %b, i32 0
+  store i64 %0, i64* %a, align 8
+  ret void
+}
+
 define void @test_vst1q_lane_f32(float* %a, <4 x float> %b) {
 ; CHECK-LABEL: test_vst1q_lane_f32:
 ; CHECK: st1 { {{v[0-9]+}}.s }[{{[0-9]+}}], [x0]

   ret void
 }
 
+define void @test_vst1q_lane0_f32(float* %a, <4 x float> %b) {
+; CHECK-LABEL: test_vst1q_lane0_f32:
+; CHECK: str {{s[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <4 x float> %b, i32 0
+  store float %0, float* %a, align 4
+  ret void
+}
+
 define void @test_vst1q_lane_f64(double* %a, <2 x double> %b) {
 ; CHECK-LABEL: test_vst1q_lane_f64:
 ; CHECK: st1 { {{v[0-9]+}}.d }[{{[0-9]+}}], [x0]

   ret void
 }
 
+define void @test_vst1q_lane0_f64(double* %a, <2 x double> %b) {
+; CHECK-LABEL: test_vst1q_lane0_f64:
+; CHECK: str {{d[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <2 x double> %b, i32 0
+  store double %0, double* %a, align 8
+  ret void
+}
+
 define void @test_vst1_lane_s8(i8* %a, <8 x i8> %b) {
 ; CHECK-LABEL: test_vst1_lane_s8:
 ; CHECK: st1 { {{v[0-9]+}}.b }[{{[0-9]+}}], [x0]

   ret void
 }
 
+define void @test_vst1_lane0_s16(i16* %a, <4 x i16> %b) {
+; CHECK-LABEL: test_vst1_lane0_s16:
+; CHECK: str {{h[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <4 x i16> %b, i32 0
+  store i16 %0, i16* %a, align 2
+  ret void
+}
+
 define void @test_vst1_lane_s32(i32* %a, <2 x i32> %b) {
 ; CHECK-LABEL: test_vst1_lane_s32:
 ; CHECK: st1 { {{v[0-9]+}}.s }[{{[0-9]+}}], [x0]

   ret void
 }
 
+define void @test_vst1_lane0_s32(i32* %a, <2 x i32> %b) {
+; CHECK-LABEL: test_vst1_lane0_s32:
+; CHECK: str {{s[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <2 x i32> %b, i32 0
+  store i32 %0, i32* %a, align 4
+  ret void
+}
+
 define void @test_vst1_lane_s64(i64* %a, <1 x i64> %b) {
 ; CHECK-LABEL: test_vst1_lane_s64:
-; CHECK: st1 { {{v[0-9]+}}.d }[{{[0-9]+}}], [x0]
+; CHECK: str {{d[0-9]+}}, [x0]
 entry:
   %0 = extractelement <1 x i64> %b, i32 0
   store i64 %0, i64* %a, align 8

   ret void
 }
 
+define void @test_vst1_lane0_f32(float* %a, <2 x float> %b) {
+; CHECK-LABEL: test_vst1_lane0_f32:
+; CHECK: str {{s[0-9]+}}, [x0]
+entry:
+  %0 = extractelement <2 x float> %b, i32 0
+  store float %0, float* %a, align 4
+  ret void
+}
+
 define void @test_vst1_lane_f64(double* %a, <1 x double> %b) {
 ; CHECK-LABEL: test_vst1_lane_f64:
 ; CHECK: str {{d[0-9]+}}, [x0]

llvm/test/CodeGen/AArch64/arm64-st1.ll

   ret void
 }
 
+define void @st1lane0_8h(<8 x i16> %A, i16* %D) {
+; CHECK-LABEL: st1lane0_8h
+; CHECK: str
+  %tmp = extractelement <8 x i16> %A, i32 0
+  store i16 %tmp, i16* %D
+  ret void
+}
+
 define void @st1lane_ro_8h(<8 x i16> %A, i16* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_8h
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_4s(<4 x i32> %A, i32* %D) {
+; CHECK-LABEL: st1lane0_4s
+; CHECK: str
+  %tmp = extractelement <4 x i32> %A, i32 0
+  store i32 %tmp, i32* %D
+  ret void
+}
+
 define void @st1lane_ro_4s(<4 x i32> %A, i32* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_4s
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_4s_float(<4 x float> %A, float* %D) {
+; CHECK-LABEL: st1lane0_4s_float
+; CHECK: str
+  %tmp = extractelement <4 x float> %A, i32 0
+  store float %tmp, float* %D
+  ret void
+}
+
 define void @st1lane_ro_4s_float(<4 x float> %A, float* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_4s_float
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_2d(<2 x i64> %A, i64* %D) {
+; CHECK-LABEL: st1lane0_2d
+; CHECK: str
+  %tmp = extractelement <2 x i64> %A, i32 0
+  store i64 %tmp, i64* %D
+  ret void
+}
+
 define void @st1lane_ro_2d(<2 x i64> %A, i64* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_2d
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_2d_double(<2 x double> %A, double* %D) {
+; CHECK-LABEL: st1lane0_2d_double
+; CHECK: str
+  %tmp = extractelement <2 x double> %A, i32 0
+  store double %tmp, double* %D
+  ret void
+}
+
 define void @st1lane_ro_2d_double(<2 x double> %A, double* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_2d_double
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_4h(<4 x i16> %A, i16* %D) {
+; CHECK-LABEL: st1lane0_4h
+; CHECK: str
+  %tmp = extractelement <4 x i16> %A, i32 0
+  store i16 %tmp, i16* %D
+  ret void
+}
+
 define void @st1lane_ro_4h(<4 x i16> %A, i16* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_4h
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_2s(<2 x i32> %A, i32* %D) {
+; CHECK-LABEL: st1lane0_2s
+; CHECK: str
+  %tmp = extractelement <2 x i32> %A, i32 0
+  store i32 %tmp, i32* %D
+  ret void
+}
+
 define void @st1lane_ro_2s(<2 x i32> %A, i32* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_2s
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

   ret void
 }
 
+define void @st1lane0_2s_float(<2 x float> %A, float* %D) {
+; CHECK-LABEL: st1lane0_2s_float
+; CHECK: str
+  %tmp = extractelement <2 x float> %A, i32 0
+  store float %tmp, float* %D
+  ret void
+}
+
 define void @st1lane_ro_2s_float(<2 x float> %A, float* %D, i64 %offset) {
 ; CHECK-LABEL: st1lane_ro_2s_float
 ; CHECK: add x[[XREG:[0-9]+]], x0, x1

llvm/test/CodeGen/AArch64/fp16-vector-load-store.ll

   ret void
 }
 
+define void @store_lane0_64(half* nocapture %a, <4 x half> %b) #1 {
+; CHECK-LABEL: store_lane0_64:
+; CHECK: str h0, [x0]
+entry:
+  %0 = extractelement <4 x half> %b, i32 0
+  store half %0, half* %a, align 2
+  ret void
+}
+
 ; Store from one lane of v8f16
 define void @store_lane_128(half* nocapture %a, <8 x half> %b) #1 {
 ; CHECK-LABEL: store_lane_128:

   ret void
 }
 
+define void @store_lane0_128(half* nocapture %a, <8 x half> %b) #1 {
+; CHECK-LABEL: store_lane0_128:
+; CHECK: str h0, [x0]
+entry:
+  %0 = extractelement <8 x half> %b, i32 0
+  store half %0, half* %a, align 2
+  ret void
+}
+
 ; NEON intrinsics - (de-)interleaving loads and stores
 declare { <4 x half>, <4 x half> } @llvm.aarch64.neon.ld2.v4f16.p0v4f16(<4 x half>*)
 declare { <4 x half>, <4 x half>, <4 x half> } @llvm.aarch64.neon.ld3.v4f16.p0v4f16(<4 x half>*)