[llvm][SVE] IR intrinscs for matrix multiplication instructions.

[fpetrogalli]

Summary

Instructions:

• SMMLA
• UMMLA
• USMMLA
• FMMLA

[fpetrogalli]

I just renamed a test (the int8 one) for consistencty.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

1864

Can you pass the intrinsic to the multiclass sve_int_matmul and add the pattenrs there? That's what is done for all the other intrinsics.

Same for the other patterns you've added.

1866

nit: _Zda_ why the double underscores?

[fpetrogalli]

Thanks for the review, @sdesmalen. I have addressed all your comments.

Francesco

[fpetrogalli]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

1866

That was copy paste coding, I have removed it. :)

[fpetrogalli]

One more _Zda_ to Zda change I missed in last patch. NFC.

[sdesmalen]

llvm/include/llvm/IR/IntrinsicsAArch64.td

2273

Can you rename this intrinsic to int_aarch64_sve_fmmla to match the instruction and make it clear the intrinsic is only valid for FP values?

llvm/lib/Target/AArch64/SVEInstrFormats.td

7553

Please use SVE_3_Op_Pat for these patterns.

7627

Please use SVE_3_Op_Pat for these patterns.

llvm/test/CodeGen/AArch64/sve-intrinsics-matmul-fp32.ll

4

nit: unnecessary whitespace :)

[fpetrogalli]

Thank you for the review, @sdesmalen!

Francesco

[fpetrogalli]

llvm/include/llvm/IR/IntrinsicsAArch64.td

2273

Yup, thank you for the heads up!

llvm/lib/Target/AArch64/SVEInstrFormats.td

7553

Neat, thanks!

[sdesmalen]

LGTM!

llvm/include/llvm/IR/IntrinsicsAArch64.td

: Intrinsic<[],
                 [
                   LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>, // Predicate
                   llvm_anyvector_ty, // Base addresses
                   llvm_i64_ty, // Scalar offset
                   llvm_i32_ty // Prfop
                 ],
                 [IntrInaccessibleMemOrArgMemOnly, ImmArg<3>]>;
 
+class SVE_MatMul_Intrinsic
+    : Intrinsic<[llvm_anyvector_ty],
+                [LLVMMatchType<0>, LLVMSubdivide4VectorType<0>, LLVMSubdivide4VectorType<0>],
+                [IntrNoMem]>;
+
 //
 // Loads
 //
 
 def int_aarch64_sve_ld1   : AdvSIMD_1Vec_PredLoad_Intrinsic;
 
 def int_aarch64_sve_ldnt1 : AdvSIMD_1Vec_PredLoad_Intrinsic;
 def int_aarch64_sve_ldnf1 : AdvSIMD_1Vec_PredLoad_Intrinsic;

 //
 // SVE2 - Optional bit permutation
 //
 
 def int_aarch64_sve_bdep_x : AdvSIMD_2VectorArg_Intrinsic;
 def int_aarch64_sve_bext_x : AdvSIMD_2VectorArg_Intrinsic;
 def int_aarch64_sve_bgrp_x : AdvSIMD_2VectorArg_Intrinsic;
 
+
+//
+// SVE ACLE: 7.3. INT8 matrix multiply extensions
+//
+def int_aarch64_sve_ummla : SVE_MatMul_Intrinsic;
+def int_aarch64_sve_smmla : SVE_MatMul_Intrinsic;
+def int_aarch64_sve_usmmla : SVE_MatMul_Intrinsic;
+
+//
+// SVE ACLE: 7.4/5. FP64/FP32 matrix multiply extensions
+//
+def int_aarch64_sve_fmmla : AdvSIMD_3VectorArg_Intrinsic;
+
 }
 
 //
 // SVE2 - Contiguous conflict detection
 //
 
 def int_aarch64_sve_whilerw_b : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilerw_h : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilerw_s : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilerw_d : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilewr_b : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilewr_h : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilewr_s : SVE2_CONFLICT_DETECT_Intrinsic;
 def int_aarch64_sve_whilewr_d : SVE2_CONFLICT_DETECT_Intrinsic;

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def : Pat<(nxv2f64 (vector_insert (nxv2f64 ZPR:$vec), (f64 FPR64:$src), GPR64:$index)),
             (CPY_ZPmV_D ZPR:$vec,
                         (CMPEQ_PPzZZ_D (PTRUE_D 31),
                                        (INDEX_II_D 0, 1),
                                        (DUP_ZR_D $index)),
                         $src)>;
 }
 
 let Predicates = [HasSVE, HasMatMulInt8] in {
-  def  SMMLA_ZZZ : sve_int_matmul<0b00, "smmla">;
-  def  UMMLA_ZZZ : sve_int_matmul<0b11, "ummla">;
-  def USMMLA_ZZZ : sve_int_matmul<0b10, "usmmla">;
+  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>;
   def USDOT_ZZZ  : sve_int_dot_mixed<"usdot">;
   def USDOT_ZZZI : sve_int_dot_mixed_indexed<0, "usdot">;
   def SUDOT_ZZZI : sve_int_dot_mixed_indexed<1, "sudot">;
 }
 
 let Predicates = [HasSVE, HasMatMulFP32] in {
-  def FMMLA_ZZZ_S : sve_fp_matrix_mla<0, "fmmla", ZPR32>;
+  defm FMMLA_ZZZ_S : sve_fp_matrix_mla<0, "fmmla", ZPR32, int_aarch64_sve_fmmla, nxv4f32>;
 }
 
 let Predicates = [HasSVE, HasMatMulFP64] in {
-  def FMMLA_ZZZ_D : sve_fp_matrix_mla<1, "fmmla", ZPR64>;
+  defm FMMLA_ZZZ_D : sve_fp_matrix_mla<1, "fmmla", ZPR64, int_aarch64_sve_fmmla, nxv2f64>;
   defm LD1RO_B_IMM  : sve_mem_ldor_si<0b00, "ld1rob", Z_b, ZPR8>;
   defm LD1RO_H_IMM  : sve_mem_ldor_si<0b01, "ld1roh", Z_h, ZPR16>;
   defm LD1RO_W_IMM  : sve_mem_ldor_si<0b10, "ld1row", Z_s, ZPR32>;
   defm LD1RO_D_IMM  : sve_mem_ldor_si<0b11, "ld1rod", Z_d, ZPR64>;
   defm LD1RO_B      : sve_mem_ldor_ss<0b00, "ld1rob", Z_b, ZPR8,  GPR64NoXZRshifted8>;
   defm LD1RO_H      : sve_mem_ldor_ss<0b01, "ld1roh", Z_h, ZPR16, GPR64NoXZRshifted16>;
   defm LD1RO_W      : sve_mem_ldor_ss<0b10, "ld1row", Z_s, ZPR32, GPR64NoXZRshifted32>;
   defm LD1RO_D      : sve_mem_ldor_ss<0b11, "ld1rod", Z_d, ZPR64, GPR64NoXZRshifted64>;

llvm/lib/Target/AArch64/SVEInstrFormats.td

: I<(outs ZPR32:$Zda), (ins ZPR32:$_Zda, ZPR8:$Zn, ZPR8:$Zm), asm,
   let Inst{9-5}   = Zn;
   let Inst{4-0}   = Zda;
 
   let Constraints = "$Zda = $_Zda";
   let DestructiveInstType = DestructiveOther;
   let ElementSize = ZPR32.ElementSize;
 }
 
+multiclass sve_int_matmul<bits<2> uns, string asm, SDPatternOperator op> {
+  def NAME : sve_int_matmul<uns, asm>;
+
+  def : SVE_3_Op_Pat<nxv4i32, op , nxv4i32, nxv16i8, nxv16i8, !cast<Instruction>(NAME)>;
+}
+
 //===----------------------------------------------------------------------===//
 // SVE Integer Dot Product Mixed Sign Group
 //===----------------------------------------------------------------------===//
 
 class sve_int_dot_mixed<string asm>
 : I<(outs ZPR32:$Zda), (ins ZPR32:$_Zda, ZPR8:$Zn, ZPR8:$Zm), asm,
   "\t$Zda, $Zn, $Zm", "", []>, Sched<[]> {
   bits<5> Zda;
: I<(outs zprty:$Zda), (ins zprty:$_Zda, zprty:$Zn, zprty:$Zm),
   let Inst{9-5}   = Zn;
   let Inst{4-0}   = Zda;
 
   let Constraints = "$Zda = $_Zda";
   let DestructiveInstType = DestructiveOther;
   let ElementSize = zprty.ElementSize;
 }
 
+multiclass sve_fp_matrix_mla<bit sz, string asm, ZPRRegOp zprty, SDPatternOperator op, ValueType vt> {
+  def NAME : sve_fp_matrix_mla<sz, asm, zprty>;
+
+  def : SVE_3_Op_Pat<vt, op , vt, vt, vt, !cast<Instruction>(NAME)>;
+}
+
 //===----------------------------------------------------------------------===//
 // SVE Memory - Contiguous Load And Replicate 256-bit Group
 //===----------------------------------------------------------------------===//
 
 class sve_mem_ldor_si<bits<2> sz, string asm, RegisterOperand VecList>
 : I<(outs VecList:$Zt), (ins PPR3bAny:$Pg, GPR64sp:$Rn, simm4s32:$imm4),
   asm, "\t$Zt, $Pg/z, [$Rn, $imm4]", "", []>, Sched<[]> {
   bits<5> Zt;

llvm/test/CodeGen/AArch64/sve-intrinsics-matmul-fp32.ll

+; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+sve,+f32mm -asm-verbose=0 < %s -o - | FileCheck %s
+
+define <vscale x 4 x float> @fmmla_s(<vscale x 4 x float> %r, <vscale x 4 x float> %a, <vscale x 4 x float> %b) nounwind {
+entry:
+; CHECK-LABEL: fmmla_s:
+; CHECK-NEXT:  fmmla   z0.s, z1.s, z2.s
+; CHECK-NEXT : ret
+  %val = tail call <vscale x 4 x float> @llvm.aarch64.sve.fmmla.nxv4f32(<vscale x 4 x float> %r, <vscale x 4 x float> %a, <vscale x 4 x float> %b)
+  ret <vscale x 4 x float> %val
+}
+
+declare <vscale x 4 x float> @llvm.aarch64.sve.fmmla.nxv4f32(<vscale x 4 x float>,<vscale x 4 x float>,<vscale x 4 x float>)
+

llvm/test/CodeGen/AArch64/sve-intrinsics-matmul-fp64.ll

+; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+sve,+f64mm -asm-verbose=0 < %s -o - | FileCheck %s
+
+
+define <vscale x 2 x double> @fmmla_d(<vscale x 2 x double> %r, <vscale x 2 x double> %a, <vscale x 2 x double> %b) nounwind {
+entry:
+; CHECK-LABEL: fmmla_d:
+; CHECK-NEXT:  fmmla   z0.d, z1.d, z2.d
+; CHECK-NEXT : ret
+  %val = tail call <vscale x 2 x double> @llvm.aarch64.sve.fmmla.nxv2f64(<vscale x 2 x double> %r, <vscale x 2 x double> %a, <vscale x 2 x double> %b)
+  ret <vscale x 2 x double> %val
+}
+
+declare <vscale x 2 x double> @llvm.aarch64.sve.fmmla.nxv2f64(<vscale x 2 x double>,<vscale x 2 x double>,<vscale x 2 x double>)
+

llvm/test/CodeGen/AArch64/sve-intrinsics-matmul-int8.ll

+; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+sve,+i8mm -asm-verbose=0 < %s -o - | FileCheck %s
+
+define <vscale x 4 x i32> @smmla(<vscale x 4 x i32> %r, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) nounwind {
+entry:
+; CHECK-LABEL: smmla:
+; CHECK-NEXT:  smmla   z0.s, z1.b, z2.b
+; CHECK-NEXT:  ret
+  %val = tail call <vscale x 4 x i32> @llvm.aarch64.sve.smmla.nxv4i32(<vscale x 4 x i32> %r, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+  ret <vscale x 4 x i32> %val
+}
+
+define <vscale x 4 x i32> @ummla(<vscale x 4 x i32> %r, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) nounwind {
+entry:
+; CHECK-LABEL: ummla:
+; CHECK-NEXT:  ummla   z0.s, z1.b, z2.b
+; CHECK-NEXT:  ret
+  %val = tail call <vscale x 4 x i32> @llvm.aarch64.sve.ummla.nxv4i32(<vscale x 4 x i32> %r, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+  ret <vscale x 4 x i32> %val
+}
+
+define <vscale x 4 x i32> @usmmla(<vscale x 4 x i32> %r, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) nounwind {
+entry:
+; CHECK-LABEL: usmmla:
+; CHECK-NEXT:  usmmla   z0.s, z1.b, z2.b
+; CHECK-NEXT : ret
+  %val = tail call <vscale x 4 x i32> @llvm.aarch64.sve.usmmla.nxv4i32(<vscale x 4 x i32> %r, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+  ret <vscale x 4 x i32> %val
+}
+
+declare <vscale x 4 x i32> @llvm.aarch64.sve.smmla.nxv4i32(<vscale x 4 x i32>, <vscale x 16 x i8>, <vscale x 16 x i8>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.ummla.nxv4i32(<vscale x 4 x i32>, <vscale x 16 x i8>, <vscale x 16 x i8>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.usmmla.nxv4i32(<vscale x 4 x i32>, <vscale x 16 x i8>, <vscale x 16 x i8>)
+