[AArch64][SVE] Mark AES instructions commutable.

[rj-jesus]

We are already doing this for the Neon versions of the instructions, just not for SVE.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Ricardo Jesus (rj-jesus)

Changes

We are already doing this for the Neon versions of the instructions, just not for SVE.

---

Full diff: https://github.com/llvm/llvm-project/pull/142919.diff

3 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td (+2-2)
• (modified) llvm/lib/Target/AArch64/SVEInstrFormats.td (+3-1)
• (modified) llvm/test/CodeGen/AArch64/sve2-intrinsics-crypto.ll (+22)

diff --git a/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td b/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td
index 91db6b6fc7984..287d0436af7c5 100644
--- a/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td
@@ -4064,8 +4064,8 @@ let Predicates = [HasSVE2_or_SME] in {
 
 let Predicates = [HasSVEAES, HasNonStreamingSVE2_or_SSVE_AES] in {
   // SVE2 crypto destructive binary operations
-  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>;
+  defm AESE_ZZZ_B : sve2_crypto_des_bin_op<0b00, "aese", ZPR8, int_aarch64_sve_aese, nxv16i8, /*commutable=*/1>;
+  defm AESD_ZZZ_B : sve2_crypto_des_bin_op<0b01, "aesd", ZPR8, int_aarch64_sve_aesd, nxv16i8, /*commutable=*/1>;
 
   // SVE2 crypto unary operations
   defm AESMC_ZZ_B  : sve2_crypto_unary_op<0b0, "aesmc",  int_aarch64_sve_aesmc>;
diff --git a/llvm/lib/Target/AArch64/SVEInstrFormats.td b/llvm/lib/Target/AArch64/SVEInstrFormats.td
index c56713783289e..d68a0bd8c7b39 100644
--- a/llvm/lib/Target/AArch64/SVEInstrFormats.td
+++ b/llvm/lib/Target/AArch64/SVEInstrFormats.td
@@ -9254,7 +9254,9 @@ class sve2_crypto_des_bin_op<bits<2> opc, string asm, ZPRRegOp zprty>
 }
 
 multiclass sve2_crypto_des_bin_op<bits<2> opc, string asm, ZPRRegOp zprty,
-                                  SDPatternOperator op, ValueType vt> {
+                                  SDPatternOperator op, ValueType vt,
+                                  bit commutable = 0> {
+  let isCommutable = commutable in
   def NAME : sve2_crypto_des_bin_op<opc, asm, zprty>;
   def : SVE_2_Op_Pat<vt, op, vt, vt, !cast<Instruction>(NAME)>;
 }
diff --git a/llvm/test/CodeGen/AArch64/sve2-intrinsics-crypto.ll b/llvm/test/CodeGen/AArch64/sve2-intrinsics-crypto.ll
index fe8271cdf054b..f477fcbe1eb5b 100644
--- a/llvm/test/CodeGen/AArch64/sve2-intrinsics-crypto.ll
+++ b/llvm/test/CodeGen/AArch64/sve2-intrinsics-crypto.ll
@@ -16,6 +16,17 @@ define <vscale x 16 x i8> @aesd_i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
   ret <vscale x 16 x i8> %out
 }
 
+define <vscale x 16 x i8> @aesd_i8_commuted(<vscale x 16 x i8> %a,
+; CHECK-LABEL: aesd_i8_commuted:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    aesd z0.b, z0.b, z1.b
+; CHECK-NEXT:    ret
+                                            <vscale x 16 x i8> %b) {
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8> %b,
+                                                        <vscale x 16 x i8> %a)
+  ret <vscale x 16 x i8> %out
+}
+
 ;
 ; AESIMC
 ;
@@ -43,6 +54,17 @@ define <vscale x 16 x i8> @aese_i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
   ret <vscale x 16 x i8> %out
 }
 
+define <vscale x 16 x i8> @aese_i8_commuted(<vscale x 16 x i8> %a,
+; CHECK-LABEL: aese_i8_commuted:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    aese z0.b, z0.b, z1.b
+; CHECK-NEXT:    ret
+                                            <vscale x 16 x i8> %b) {
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8> %b,
+                                                        <vscale x 16 x i8> %a)
+  ret <vscale x 16 x i8> %out
+}
+
 ;
 ; AESMC
 ;

[rj-jesus]

Is there a cleaner/preferable way of doing this?

[paulwalker-arm]

Can you just wrap the two instruction definitions in let isCommutable = 1 in {}? or is the pattern matching def preventing that?

[rj-jesus]

I tried that first, but the pattern matching gets in the way.
We can't set isCommutable unconditionally in sve2_crypto_des_bin_op either (just for the instruction def) because SM4E uses the same multiclass.
Is it worth creating a dedicated multiclass just for AES instructions as Neon seems to have?

[paulwalker-arm]

No, worst come to worst I'll be happy with the current approach.

Just before that though... The multiclass only has a single pattern and we're unlikely to add more so I think you can remove the multiclass and change the main class's DAG from [] to [(set (vt zprty:$Zdn), (op (vt zprty:$_Zdn), (vt zprty:$Zm)))] and then you'll be able to wrap the relevant definitions with isCommutable.

[rj-jesus]

Ah, thanks very much, that looks much better. :)
(Please let me know if I missed or misunderstood anything.)