x86: add VAES instructions

The AES-NI AVX instructions are extended to 256-bit operands. Fixes #630.
jasmin-lang · Jun 10, 2024 · 4a4938c · 4a4938c
1 parent 5249c4e
commit 4a4938c
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 20 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,11 @@
 
 ## New features
 
+- The instructions of the VAES extension are available through a size suffix
+  (e.g., `VAESENC_256`)
+  ([PR #831](https://github.com/jasmin-lang/jasmin/pull/831),
+  fixes [#630](https://github.com/jasmin-lang/jasmin/issues/630)).
+
 - The following x86 BMI2 instructions are now available:
   `RORX`, `SARX`, `SHRX`, and `SHLX`
   ([PR #824](https://github.com/jasmin-lang/jasmin/pull/824)).

diff --git a/compiler/tests/success/x86-64/aes_instr.jazz b/compiler/tests/success/x86-64/aes_instr.jazz
@@ -42,4 +42,19 @@ x1 = #VAESKEYGENASSIST((u128)[p], 0);
 
 r = 0;
 return r;
-}
+}
+
+export
+fn test_vaes(reg u64 p) {
+  reg u256 x y z;
+  x = (u256)[p];
+  y = #VAESENC_256(x, (u256)[p + 32]);
+  z = #VAESENC_256(y, x);
+  x = #VAESENCLAST_256(z, (u256)[p + 64]);
+  y = #VAESENCLAST_256(x, y);
+  z = #VAESDEC_256(y, (u256)[p + 32]);
+  x = #VAESDEC_256(z, y);
+  y = #VAESDECLAST_256(x, (u256)[p]);
+  z = #VAESDECLAST_256(y, x);
+  (u256)[p] = z;
+}
diff --git a/eclib/JModel_x86.ec b/eclib/JModel_x86.ec
@@ -923,12 +923,21 @@ op VPTEST_256 (v1 v2: W256.t) =
 (* AES instruction *)
 
 abbrev [-printing] VAESDEC          = AESDEC.
+abbrev [-printing] VAESDEC_128      = AESDEC.
 abbrev [-printing] VAESDECLAST      = AESDECLAST.
+abbrev [-printing] VAESDECLAST_128  = AESDECLAST.
 abbrev [-printing] VAESENC          = AESENC.
+abbrev [-printing] VAESENC_128      = AESENC.
 abbrev [-printing] VAESENCLAST      = AESENCLAST.
+abbrev [-printing] VAESENCLAST_128  = AESENCLAST.
 abbrev [-printing] VAESIMC          = AESIMC.
 abbrev [-printing] VAESKEYGENASSIST = AESKEYGENASSIST.
 
+op VAESDEC_256     = W2u128.map2 VAESDEC_128.
+op VAESDECLAST_256 = W2u128.map2 VAESDECLAST_128.
+op VAESENC_256     = W2u128.map2 VAESENC_128.
+op VAESENCLAST_256 = W2u128.map2 VAESENCLAST_128.
+
 (* ------------------------------------------------------------------- *)
 (* PCLMULQDQ instructions *)
 (*

diff --git a/proofs/compiler/x86_instr_decl.v b/proofs/compiler/x86_instr_decl.v
@@ -168,13 +168,13 @@ Variant x86_op : Type :=
 
 (* AES instructions *)
 | AESDEC
-| VAESDEC
+| VAESDEC of wsize
 | AESDECLAST
-| VAESDECLAST
+| VAESDECLAST of wsize
 | AESENC
-| VAESENC
+| VAESENC of wsize
 | AESENCLAST
-| VAESENCLAST
+| VAESENCLAST of wsize
 | AESIMC
 | VAESIMC
 | AESKEYGENASSIST
@@ -1873,33 +1873,37 @@ Definition mk_instr_aes2 jname aname (constr:x86_op) x86_sem msb_flag :=
   mk_instr_pp jname (w2_ty U128 U128) (w_ty U128) [:: Eu 0; Eu 1] [:: Eu 0] msb_flag x86_sem
          (check_xmm_xmmm U128) 2 (primM constr) (pp_name_ty aname [::U128;U128]).
 
-Definition mk_instr_aes3 jname aname (constr:x86_op) x86_sem msb_flag :=
-  mk_instr_pp jname (w2_ty U128 U128) (w_ty U128) [:: Eu 1; Eu 2] [:: Eu 0] msb_flag x86_sem
-         (check_xmm_xmm_xmmm U128) 3 (primM constr) (pp_name_ty aname [::U128;U128;U128]).
+Definition mk_instr_aes3 jname aname (constr: wsize → x86_op) x86_sem :=
+  (λ sz, mk_instr (pp_sz jname sz) (w2_ty sz sz) (w_ty sz) [:: Eu 1; Eu 2]
+           [:: Eu 0] MSB_CLEAR
+           (x86_u128_binop (lift2_vec U128 x86_sem sz))
+           (check_xmm_xmm_xmmm sz) 3 [::]
+           (pp_name_ty aname [:: sz; sz; sz ]),
+   (jname%string, prim_128_256 constr)).
 
 Definition Ox86_AESDEC_instr := 
   mk_instr_aes2 "AESDEC" "aesdec" AESDEC x86_AESDEC MSB_MERGE.
 
-Definition Ox86_VAESDEC_instr := 
-  mk_instr_aes3 "VAESDEC" "vaesdec" VAESDEC x86_AESDEC MSB_CLEAR.
+Definition Ox86_VAESDEC_instr :=
+  mk_instr_aes3 "VAESDEC" "vaesdec" VAESDEC wAESDEC.
 
 Definition Ox86_AESDECLAST_instr := 
   mk_instr_aes2 "AESDECLAST" "aesdeclast" AESDECLAST x86_AESDECLAST MSB_MERGE.
 
-Definition Ox86_VAESDECLAST_instr := 
-  mk_instr_aes3 "VAESDECLAST" "vaesdeclast" VAESDECLAST x86_AESDECLAST MSB_CLEAR.
+Definition Ox86_VAESDECLAST_instr :=
+  mk_instr_aes3 "VAESDECLAST" "vaesdeclast" VAESDECLAST wAESDECLAST.
 
 Definition Ox86_AESENC_instr := 
   mk_instr_aes2 "AESENC" "aesenc" AESENC x86_AESENC MSB_MERGE.
 
-Definition Ox86_VAESENC_instr := 
-  mk_instr_aes3 "VAESENC" "vaesenc" VAESENC x86_AESENC MSB_CLEAR.
+Definition Ox86_VAESENC_instr :=
+  mk_instr_aes3 "VAESENC" "vaesenc" VAESENC wAESENC.
 
 Definition Ox86_AESENCLAST_instr := 
   mk_instr_aes2 "AESENCLAST" "aesenclast" AESENCLAST x86_AESENCLAST MSB_MERGE.
 
-Definition Ox86_VAESENCLAST_instr := 
-  mk_instr_aes3 "VAESENCLAST" "vaesenclast" VAESENCLAST x86_AESENCLAST MSB_CLEAR.
+Definition Ox86_VAESENCLAST_instr :=
+  mk_instr_aes3 "VAESENCLAST" "vaesenclast" VAESENCLAST wAESENCLAST.
 
 Definition Ox86_AESIMC_instr := 
   mk_instr_pp "AESIMC" (w_ty U128) (w_ty U128) [:: Eu 1] [:: Eu 0] MSB_MERGE x86_AESIMC
@@ -2063,13 +2067,13 @@ Definition x86_instr_desc o : instr_desc_t :=
   | RDTSC sz           => Ox86_RDTSC_instr.1 sz
   | RDTSCP sz          => Ox86_RDTSCP_instr.1 sz
   | AESDEC             => Ox86_AESDEC_instr.1          
-  | VAESDEC            => Ox86_VAESDEC_instr.1         
+  | VAESDEC sz         => Ox86_VAESDEC_instr.1 sz
   | AESDECLAST         => Ox86_AESDECLAST_instr.1      
-  | VAESDECLAST        => Ox86_VAESDECLAST_instr.1     
+  | VAESDECLAST sz     => Ox86_VAESDECLAST_instr.1 sz
   | AESENC             => Ox86_AESENC_instr.1          
-  | VAESENC            => Ox86_VAESENC_instr.1         
+  | VAESENC sz         => Ox86_VAESENC_instr.1 sz
   | AESENCLAST         => Ox86_AESENCLAST_instr.1      
-  | VAESENCLAST        => Ox86_VAESENCLAST_instr.1     
+  | VAESENCLAST sz     => Ox86_VAESENCLAST_instr.1 sz
   | AESIMC             => Ox86_AESIMC_instr.1          
   | VAESIMC            => Ox86_VAESIMC_instr.1         
   | AESKEYGENASSIST    => Ox86_AESKEYGENASSIST_instr.1