[AArch64][SVE] Add intrinsics for SVE2 cryptographic instructions

[kmclaughlin]

Summary

Implements the following SVE2 intrinsics:

• @llvm.aarch64.sve.aesd
• @llvm.aarch64.sve.aesimc
• @llvm.aarch64.sve.aese
• @llvm.aarch64.sve.aesmc
• @llvm.aarch64.sve.rax1
• @llvm.aarch64.sve.sm4e
• @llvm.aarch64.sve.sm4ekey

[sdesmalen]

LGTM!

llvm/include/llvm/IR/IntrinsicsAArch64.td

 // SVE2 - Polynomial arithmetic
 //
 
 def int_aarch64_sve_eorbt       : AdvSIMD_3VectorArg_Intrinsic;
 def int_aarch64_sve_eortb       : AdvSIMD_3VectorArg_Intrinsic;
 def int_aarch64_sve_pmullb_pair : AdvSIMD_2VectorArg_Intrinsic;
 def int_aarch64_sve_pmullt_pair : AdvSIMD_2VectorArg_Intrinsic;
 
+//
+// SVE2 - Optional AES, SHA-3 and SM4
+//
+
+def int_aarch64_sve_aesd    : GCCBuiltin<"__builtin_sve_svaesd_u8">,
+                              Intrinsic<[llvm_nxv16i8_ty],
+                                        [llvm_nxv16i8_ty, llvm_nxv16i8_ty],
+                                        [IntrNoMem]>;
+def int_aarch64_sve_aesimc  : GCCBuiltin<"__builtin_sve_svaesimc_u8">,
+                              Intrinsic<[llvm_nxv16i8_ty],
+                                        [llvm_nxv16i8_ty],
+                                        [IntrNoMem]>;
+def int_aarch64_sve_aese    : GCCBuiltin<"__builtin_sve_svaese_u8">,
+                              Intrinsic<[llvm_nxv16i8_ty],
+                                        [llvm_nxv16i8_ty, llvm_nxv16i8_ty],
+                                        [IntrNoMem]>;
+def int_aarch64_sve_aesmc   : GCCBuiltin<"__builtin_sve_svaesmc_u8">,
+                              Intrinsic<[llvm_nxv16i8_ty],
+                                        [llvm_nxv16i8_ty],
+                                        [IntrNoMem]>;
+def int_aarch64_sve_rax1    : GCCBuiltin<"__builtin_sve_svrax1_u64">,
+                              Intrinsic<[llvm_nxv2i64_ty],
+                                        [llvm_nxv2i64_ty, llvm_nxv2i64_ty],
+                                        [IntrNoMem]>;
+def int_aarch64_sve_sm4e    : GCCBuiltin<"__builtin_sve_svsm4e_u32">,
+                              Intrinsic<[llvm_nxv4i32_ty],
+                                        [llvm_nxv4i32_ty, llvm_nxv4i32_ty],
+                                        [IntrNoMem]>;
+def int_aarch64_sve_sm4ekey : GCCBuiltin<"__builtin_sve_svsm4ekey_u32">,
+                              Intrinsic<[llvm_nxv4i32_ty],
+                                        [llvm_nxv4i32_ty, llvm_nxv4i32_ty],
+                                        [IntrNoMem]>;
+
 }

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVE2] in {
 
   // SVE2 pointer conflict compare
   defm WHILEWR_PXX : sve2_int_while_rr<0b0, "whilewr">;
   defm WHILERW_PXX : sve2_int_while_rr<0b1, "whilerw">;
 }
 
 let Predicates = [HasSVE2AES] in {
   // SVE2 crypto destructive binary operations
-  def AESE_ZZZ_B : sve2_crypto_des_bin_op<0b00, "aese", ZPR8>;
-  def AESD_ZZZ_B : sve2_crypto_des_bin_op<0b01, "aesd", ZPR8>;
+  defm AESE_ZZZ_B : sve2_crypto_des_bin_op<0b00, "aese", ZPR8, int_aarch64_sve_aese, nxv16i8>;
+  defm AESD_ZZZ_B : sve2_crypto_des_bin_op<0b01, "aesd", ZPR8, int_aarch64_sve_aesd, nxv16i8>;
 
   // SVE2 crypto unary operations
-  def AESMC_ZZ_B  : sve2_crypto_unary_op<0b0, "aesmc">;
-  def AESIMC_ZZ_B : sve2_crypto_unary_op<0b1, "aesimc">;
+  defm AESMC_ZZ_B  : sve2_crypto_unary_op<0b0, "aesmc",  int_aarch64_sve_aesmc>;
+  defm AESIMC_ZZ_B : sve2_crypto_unary_op<0b1, "aesimc", int_aarch64_sve_aesimc>;
 
   // PMULLB and PMULLT instructions which operate with 64-bit source and
   // 128-bit destination elements are enabled with crypto extensions, similar
   // to NEON PMULL2 instruction.
   defm PMULLB_ZZZ_Q : sve2_wide_int_arith_pmul<0b00, 0b11010, "pmullb", int_aarch64_sve_pmullb_pair>;
   defm PMULLT_ZZZ_Q : sve2_wide_int_arith_pmul<0b00, 0b11011, "pmullt", int_aarch64_sve_pmullt_pair>;
 }
 
 let Predicates = [HasSVE2SM4] in {
   // SVE2 crypto constructive binary operations
-  def SM4EKEY_ZZZ_S : sve2_crypto_cons_bin_op<0b0, "sm4ekey", ZPR32>;
+  defm SM4EKEY_ZZZ_S : sve2_crypto_cons_bin_op<0b0, "sm4ekey", ZPR32, int_aarch64_sve_sm4ekey, nxv4i32>;
   // SVE2 crypto destructive binary operations
-  def SM4E_ZZZ_S : sve2_crypto_des_bin_op<0b10, "sm4e", ZPR32>;
+  defm SM4E_ZZZ_S : sve2_crypto_des_bin_op<0b10, "sm4e", ZPR32, int_aarch64_sve_sm4e, nxv4i32>;
 }
 
 let Predicates = [HasSVE2SHA3] in {
   // SVE2 crypto constructive binary operations
-  def RAX1_ZZZ_D : sve2_crypto_cons_bin_op<0b1, "rax1",    ZPR64>;
+  defm RAX1_ZZZ_D : sve2_crypto_cons_bin_op<0b1, "rax1", ZPR64, int_aarch64_sve_rax1, nxv2i64>;
 }
 
 let Predicates = [HasSVE2BitPerm] in {
   // SVE2 bitwise permute
   defm BEXT_ZZZ : sve2_misc_bitwise<0b1100, "bext">;
   defm BDEP_ZZZ : sve2_misc_bitwise<0b1101, "bdep">;
   defm BGRP_ZZZ : sve2_misc_bitwise<0b1110, "bgrp">;
 }

llvm/lib/Target/AArch64/SVEInstrFormats.td

: I<(outs zprty:$Zd), (ins zprty:$Zn, zprty:$Zm),
   let Inst{31-21} = 0b01000101001;
   let Inst{20-16} = Zm;
   let Inst{15-11} = 0b11110;
   let Inst{10}    = opc;
   let Inst{9-5}   = Zn;
   let Inst{4-0}   = Zd;
 }
 
+multiclass sve2_crypto_cons_bin_op<bit opc, string asm, ZPRRegOp zprty,
+                                   SDPatternOperator op, ValueType vt> {
+  def NAME : sve2_crypto_cons_bin_op<opc, asm, zprty>;
+  def : SVE_2_Op_Pat<vt, op, vt, vt, !cast<Instruction>(NAME)>;
+}
+
 class sve2_crypto_des_bin_op<bits<2> opc, string asm, ZPRRegOp zprty>
 : I<(outs zprty:$Zdn), (ins zprty:$_Zdn, zprty:$Zm),
   asm, "\t$Zdn, $_Zdn, $Zm",
   "",
   []>, Sched<[]> {
   bits<5> Zdn;
   bits<5> Zm;
   let Inst{31-17} = 0b010001010010001;
   let Inst{16}    = opc{1};
   let Inst{15-11} = 0b11100;
   let Inst{10}    = opc{0};
   let Inst{9-5}   = Zm;
   let Inst{4-0}   = Zdn;
 
   let Constraints = "$Zdn = $_Zdn";
 }
 
-class sve2_crypto_unary_op<bit opc, string asm>
-: I<(outs ZPR8:$Zdn), (ins ZPR8:$_Zdn),
+multiclass sve2_crypto_des_bin_op<bits<2> opc, string asm, ZPRRegOp zprty,
+                                  SDPatternOperator op, ValueType vt> {
+  def NAME : sve2_crypto_des_bin_op<opc, asm, zprty>;
+  def : SVE_2_Op_Pat<vt, op, vt, vt, !cast<Instruction>(NAME)>;
+}
+
+class sve2_crypto_unary_op<bit opc, string asm, ZPRRegOp zprty>
+: I<(outs zprty:$Zdn), (ins zprty:$_Zdn),
   asm, "\t$Zdn, $_Zdn",
   "",
   []>, Sched<[]> {
   bits<5> Zdn;
   let Inst{31-11} = 0b010001010010000011100;
   let Inst{10}    = opc;
   let Inst{9-5}   = 0b00000;
   let Inst{4-0}   = Zdn;
 
   let Constraints = "$Zdn = $_Zdn";
 }
+
+multiclass sve2_crypto_unary_op<bit opc, string asm, SDPatternOperator op> {
+  def NAME : sve2_crypto_unary_op<opc, asm, ZPR8>;
+  def : SVE_1_Op_Pat<nxv16i8, op, nxv16i8, !cast<Instruction>(NAME)>;
+}

llvm/test/CodeGen/AArch64/sve2-intrinsics-crypto.ll

+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve2-aes,+sve2-sha3,+sve2-sm4 -asm-verbose=0 < %s | FileCheck %s
+
+;
+; AESD
+;
+
+define <vscale x 16 x i8> @aesd_i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; CHECK-LABEL: aesd_i8:
+; CHECK: aesd z0.b, z0.b, z1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8> %a,
+                                                        <vscale x 16 x i8> %b)
+  ret <vscale x 16 x i8> %out
+}
+
+;
+; AESIMC
+;
+
+define <vscale x 16 x i8> @aesimc_i8(<vscale x 16 x i8> %a) {
+; CHECK-LABEL: aesimc_i8:
+; CHECK: aesimc z0.b, z0.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.aesimc(<vscale x 16 x i8> %a)
+  ret <vscale x 16 x i8> %out
+}
+
+;
+; AESE
+;
+
+define <vscale x 16 x i8> @aese_i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; CHECK-LABEL: aese_i8:
+; CHECK: aese z0.b, z0.b, z1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8> %a,
+                                                        <vscale x 16 x i8> %b)
+  ret <vscale x 16 x i8> %out
+}
+
+;
+; AESMC
+;
+
+define <vscale x 16 x i8> @aesmc_i8(<vscale x 16 x i8> %a) {
+; CHECK-LABEL: aesmc_i8:
+; CHECK: aesmc z0.b, z0.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.aesmc(<vscale x 16 x i8> %a)
+  ret <vscale x 16 x i8> %out
+}
+
+;
+; RAX1
+;
+
+define <vscale x 2 x i64> @rax1_i64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) {
+; CHECK-LABEL: rax1_i64:
+; CHECK: rax1 z0.d, z0.d, z1.d
+; CHECK-NEXT: ret
+  %out = call <vscale x 2 x i64> @llvm.aarch64.sve.rax1(<vscale x 2 x i64> %a,
+                                                        <vscale x 2 x i64> %b)
+  ret <vscale x 2 x i64> %out
+}
+
+;
+; SM4E
+;
+
+define <vscale x 4 x i32> @sm4e_i32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: sm4e_i32:
+; CHECK: sm4e z0.s, z0.s, z1.s
+; CHECK-NEXT: ret
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.sm4e(<vscale x 4 x i32> %a,
+                                                        <vscale x 4 x i32> %b)
+  ret <vscale x 4 x i32> %out
+}
+
+;
+; SM4EKEY
+;
+
+define <vscale x 4 x i32> @sm4ekey_i32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: sm4ekey_i32:
+; CHECK: sm4ekey z0.s, z0.s, z1.s
+; CHECK-NEXT: ret
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.sm4ekey(<vscale x 4 x i32> %a,
+                                                           <vscale x 4 x i32> %b)
+  ret <vscale x 4 x i32> %out
+}
+
+
+declare <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8>, <vscale x 16 x i8>)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.aesimc(<vscale x 16 x i8>)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8>, <vscale x 16 x i8>)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.aesmc(<vscale x 16 x i8>)
+declare <vscale x 2 x i64> @llvm.aarch64.sve.rax1(<vscale x 2 x i64>, <vscale x 2 x i64>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sm4e(<vscale x 4 x i32>, <vscale x 4 x i32>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sm4ekey(<vscale x 4 x i32>, <vscale x 4 x i32>)