[AArch64][SVE2p1] Add 2-way SVE2p1 dot product intrinsics

[david-arm]

Summary

This patch adds the LLVM IR intrinsics for the following:

• sdot (2-way, vectors + indexed)
• udot (2-way, vectors + indexed)
• fdot (vectors + indexed)

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {
   def int_aarch64_sve_zip_x2  : SVE2_VG2_ZipUzp_Intrinsic;
   def int_aarch64_sve_zipq_x2 : SVE2_VG2_ZipUzp_Intrinsic;
   def int_aarch64_sve_zip_x4  : SVE2_VG4_ZipUzp_Intrinsic;
   def int_aarch64_sve_zipq_x4 : SVE2_VG4_ZipUzp_Intrinsic;
   def int_aarch64_sve_uzp_x2  : SVE2_VG2_ZipUzp_Intrinsic;
   def int_aarch64_sve_uzpq_x2 : SVE2_VG2_ZipUzp_Intrinsic;
   def int_aarch64_sve_uzp_x4  : SVE2_VG4_ZipUzp_Intrinsic;
   def int_aarch64_sve_uzpq_x4 : SVE2_VG4_ZipUzp_Intrinsic;
+
+  // Vector dot-products (2-way)
+  def int_aarch64_sve_sdot_x2 : SVE2_3VectorArg_Long_Intrinsic;
+  def int_aarch64_sve_udot_x2 : SVE2_3VectorArg_Long_Intrinsic;
+  def int_aarch64_sve_fdot_x2 : SVE2_3VectorArg_Long_Intrinsic;
+  def int_aarch64_sve_sdot_lane_x2 : SVE2_3VectorArgIndexed_Long_Intrinsic;
+  def int_aarch64_sve_udot_lane_x2 : SVE2_3VectorArgIndexed_Long_Intrinsic;
+  def int_aarch64_sve_fdot_lane_x2 : SVE2_3VectorArgIndexed_Long_Intrinsic;
 }

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

 } // End HasSVE2p1_or_HasSME
 
 //===----------------------------------------------------------------------===//
 // SME2 or SVE2.1 instructions
 //===----------------------------------------------------------------------===//
 
 let Predicates = [HasSVE2p1_or_HasSME2] in {
 defm FCLAMP_ZZZ : sve2p1_fclamp<"fclamp", int_aarch64_sve_fclamp>;
-def FDOT_ZZZ_S  : sve_float_dot<0b0, "fdot">;
-def FDOT_ZZZI_S : sve_float_dot_indexed<0b0, "fdot">;
+defm FDOT_ZZZ_S  : sve_float_dot<0b0, "fdot", nxv8f16, int_aarch64_sve_fdot_x2>;
+defm FDOT_ZZZI_S : sve_float_dot_indexed<0b0, "fdot", nxv8f16, int_aarch64_sve_fdot_lane_x2>;
 def BFMLSLB_ZZZ_S : sve2_fp_mla_long<0b110, "bfmlslb">;
 def BFMLSLT_ZZZ_S : sve2_fp_mla_long<0b111, "bfmlslt">;
 def BFMLSLB_ZZZI_S : sve2_fp_mla_long_by_indexed_elem<0b110, "bfmlslb">;
 def BFMLSLT_ZZZI_S : sve2_fp_mla_long_by_indexed_elem<0b111, "bfmlslt">;
 
-def SDOT_ZZZ_HtoS  : sve2p1_two_way_dot_vv<"sdot", 0b0>;
-def UDOT_ZZZ_HtoS  : sve2p1_two_way_dot_vv<"udot", 0b1>;
-def SDOT_ZZZI_HtoS : sve2p1_two_way_dot_vvi<"sdot", 0b0>;
-def UDOT_ZZZI_HtoS : sve2p1_two_way_dot_vvi<"udot", 0b1>;
+defm SDOT_ZZZ_HtoS  : sve2p1_two_way_dot_vv<"sdot", 0b0, int_aarch64_sve_sdot_x2>;
+defm UDOT_ZZZ_HtoS  : sve2p1_two_way_dot_vv<"udot", 0b1, int_aarch64_sve_udot_x2>;
+defm SDOT_ZZZI_HtoS : sve2p1_two_way_dot_vvi<"sdot", 0b0, int_aarch64_sve_sdot_lane_x2>;
+defm UDOT_ZZZI_HtoS : sve2p1_two_way_dot_vvi<"udot", 0b1, int_aarch64_sve_udot_lane_x2>;
 
 defm CNTP_XCI : sve2p1_pcount_pn<"cntp", 0b000>;
 defm PEXT_PCI : sve2p1_pred_as_ctr_to_mask<"pext">;
 defm PEXT_2PCI : sve2p1_pred_as_ctr_to_mask_pair<"pext">;
 defm PTRUE_C  : sve2p1_ptrue_pn<"ptrue">;
 
 defm SQCVTN_Z2Z_StoH  : sve2p1_multi_vec_extract_narrow<"sqcvtn", 0b00, int_aarch64_sve_sqcvtn_x2>;
 defm UQCVTN_Z2Z_StoH  : sve2p1_multi_vec_extract_narrow<"uqcvtn", 0b01, int_aarch64_sve_uqcvtn_x2>;

llvm/lib/Target/AArch64/SVEInstrFormats.td

class sve2p1_two_way_dot_vv<string mnemonic, bit u>
   let Inst{9-5}   = Zn;
   let Inst{4-0}   = Zda;
 
   let Constraints = "$Zda = $_Zda";
   let DestructiveInstType = DestructiveOther;
   let hasSideEffects = 0;
 }
 
+multiclass sve2p1_two_way_dot_vv<string mnemonic, bit u, SDPatternOperator intrinsic> {
+  def NAME : sve2p1_two_way_dot_vv<mnemonic, u>;
+
+  def : SVE_3_Op_Pat<nxv4i32, intrinsic, nxv4i32, nxv8i16, nxv8i16, !cast<Instruction>(NAME)>;
+}
 
 // SVE two-way dot product (indexed)
 class sve2p1_two_way_dot_vvi<string mnemonic, bit u>
     : I<(outs ZPR32:$Zda), (ins ZPR32:$_Zda, ZPR16:$Zn, ZPR3b16:$Zm, VectorIndexS32b:$i2),
         mnemonic, "\t$Zda, $Zn, $Zm$i2",
         "", []>, Sched<[]> {
   bits<5> Zda;
   bits<5> Zn;
   bits<3> Zm;
   bits<2> i2;
   let Inst{31-21} = 0b01000100100;
   let Inst{20-19} = i2;
   let Inst{18-16} = Zm;
   let Inst{15-11} = 0b11001;
   let Inst{10}    = u;
   let Inst{9-5}   = Zn;
   let Inst{4-0}   = Zda;
 
   let Constraints = "$Zda = $_Zda";
   let DestructiveInstType = DestructiveOther;
   let hasSideEffects = 0;
 }
 
+multiclass sve2p1_two_way_dot_vvi<string mnemonic, bit u, SDPatternOperator intrinsic> {
+  def NAME : sve2p1_two_way_dot_vvi<mnemonic, u>;
+
+  def : SVE_4_Op_Imm_Pat<nxv4i32, intrinsic, nxv4i32, nxv8i16, nxv8i16, i32, VectorIndexS32b_timm, !cast<Instruction>(NAME)>;
+}
 
 class sve2p1_ptrue_pn<string mnemonic, bits<2> sz, PNRP8to15RegOp pnrty>
     : I<(outs pnrty:$PNd), (ins ), mnemonic, "\t$PNd",
         "", []>, Sched<[]> {
   bits<3> PNd;
   let Inst{31-24}  = 0b00100101;
   let Inst{23-22} = sz;
   let Inst{21-3}  = 0b1000000111100000010;

llvm/test/CodeGen/AArch64/sve2p1-intrinsics-dots.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve2p1 < %s | FileCheck %s
+
+define <vscale x 4 x i32> @sdot_x2(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm) {
+; CHECK-LABEL: sdot_x2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sdot z0.s, z1.h, z2.h
+; CHECK-NEXT:    ret
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.sdot.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm)
+  ret <vscale x 4 x i32> %out
+}
+
+define <vscale x 4 x i32> @udot_x2(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm) {
+; CHECK-LABEL: udot_x2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    udot z0.s, z1.h, z2.h
+; CHECK-NEXT:    ret
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.udot.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm)
+  ret <vscale x 4 x i32> %out
+}
+
+define <vscale x 4 x float> @fdot_x2(<vscale x 4 x float> %zda, <vscale x 8 x half> %zn, <vscale x 8 x half> %zm) {
+; CHECK-LABEL: fdot_x2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fdot z0.s, z1.h, z2.h
+; CHECK-NEXT:    ret
+  %out = call <vscale x 4 x float> @llvm.aarch64.sve.fdot.x2.nxv4f32(<vscale x 4 x float> %zda, <vscale x 8 x half> %zn, <vscale x 8 x half> %zm)
+  ret <vscale x 4 x float> %out
+}
+
+define <vscale x 4 x i32> @sdot_lane_x2(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm) {
+; CHECK-LABEL: sdot_lane_x2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sdot z0.s, z1.h, z2.h[3]
+; CHECK-NEXT:    ret
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.sdot.lane.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm, i32 3)
+  ret <vscale x 4 x i32> %out
+}
+
+define <vscale x 4 x i32> @udot_lane_x2(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm) {
+; CHECK-LABEL: udot_lane_x2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    udot z0.s, z1.h, z2.h[3]
+; CHECK-NEXT:    ret
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.udot.lane.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm, i32 3)
+  ret <vscale x 4 x i32> %out
+}
+
+define <vscale x 4 x float> @fdot_lane_x2(<vscale x 4 x float> %zda, <vscale x 8 x half> %zn, <vscale x 8 x half> %zm) {
+; CHECK-LABEL: fdot_lane_x2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fdot z0.s, z1.h, z2.h[3]
+; CHECK-NEXT:    ret
+  %out = call <vscale x 4 x float> @llvm.aarch64.sve.fdot.lane.x2.nxv4f32(<vscale x 4 x float> %zda, <vscale x 8 x half> %zn, <vscale x 8 x half> %zm, i32 3)
+  ret <vscale x 4 x float> %out
+}
+
+
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sdot.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.udot.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm)
+declare <vscale x 4 x float> @llvm.aarch64.sve.fdot.x2.nxv4f32(<vscale x 4 x float> %zda, <vscale x 8 x half> %zn, <vscale x 8 x half> %zm)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sdot.lane.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm, i32)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.udot.lane.x2.nxv4i32(<vscale x 4 x i32> %zda, <vscale x 8 x i16> %zn, <vscale x 8 x i16> %zm, i32)
+declare <vscale x 4 x float> @llvm.aarch64.sve.fdot.lane.x2.nxv4f32(<vscale x 4 x float> %zda, <vscale x 8 x half> %zn, <vscale x 8 x half> %zm, i32)