Skip to content
This repository was archived by the owner on May 28, 2025. It is now read-only.

Commit aaba0a5

Browse files
committed
Auto merge of rust-lang#3101 - eduardosm:x86-aes-intrinsics, r=RalfJung
Implement `llvm.x86.aesni.*` intrinsics
2 parents ec1db97 + 700bc97 commit aaba0a5

File tree

5 files changed

+539
-0
lines changed

5 files changed

+539
-0
lines changed

src/tools/miri/Cargo.lock

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,17 @@ version = "1.0.2"
1717
source = "registry+https://github.com/rust-lang/crates.io-index"
1818
checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
1919

20+
[[package]]
21+
name = "aes"
22+
version = "0.8.3"
23+
source = "registry+https://github.com/rust-lang/crates.io-index"
24+
checksum = "ac1f845298e95f983ff1944b728ae08b8cebab80d684f0a832ed0fc74dfa27e2"
25+
dependencies = [
26+
"cfg-if",
27+
"cipher",
28+
"cpufeatures",
29+
]
30+
2031
[[package]]
2132
name = "aho-corasick"
2233
version = "1.1.1"
@@ -142,6 +153,16 @@ version = "1.0.0"
142153
source = "registry+https://github.com/rust-lang/crates.io-index"
143154
checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
144155

156+
[[package]]
157+
name = "cipher"
158+
version = "0.4.4"
159+
source = "registry+https://github.com/rust-lang/crates.io-index"
160+
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
161+
dependencies = [
162+
"crypto-common",
163+
"inout",
164+
]
165+
145166
[[package]]
146167
name = "color-eyre"
147168
version = "0.6.2"
@@ -199,6 +220,15 @@ dependencies = [
199220
"windows-sys 0.45.0",
200221
]
201222

223+
[[package]]
224+
name = "cpufeatures"
225+
version = "0.2.9"
226+
source = "registry+https://github.com/rust-lang/crates.io-index"
227+
checksum = "a17b76ff3a4162b0b27f354a0c87015ddad39d35f9c0c36607a3bdd175dde1f1"
228+
dependencies = [
229+
"libc",
230+
]
231+
202232
[[package]]
203233
name = "crossbeam-channel"
204234
version = "0.5.8"
@@ -218,6 +248,16 @@ dependencies = [
218248
"cfg-if",
219249
]
220250

251+
[[package]]
252+
name = "crypto-common"
253+
version = "0.1.6"
254+
source = "registry+https://github.com/rust-lang/crates.io-index"
255+
checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3"
256+
dependencies = [
257+
"generic-array",
258+
"typenum",
259+
]
260+
221261
[[package]]
222262
name = "ctrlc"
223263
version = "3.4.1"
@@ -284,6 +324,16 @@ version = "2.0.1"
284324
source = "registry+https://github.com/rust-lang/crates.io-index"
285325
checksum = "25cbce373ec4653f1a01a31e8a5e5ec0c622dc27ff9c4e6606eefef5cbbed4a5"
286326

327+
[[package]]
328+
name = "generic-array"
329+
version = "0.14.7"
330+
source = "registry+https://github.com/rust-lang/crates.io-index"
331+
checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
332+
dependencies = [
333+
"typenum",
334+
"version_check",
335+
]
336+
287337
[[package]]
288338
name = "getrandom"
289339
version = "0.2.10"
@@ -332,6 +382,15 @@ dependencies = [
332382
"unicode-width",
333383
]
334384

385+
[[package]]
386+
name = "inout"
387+
version = "0.1.3"
388+
source = "registry+https://github.com/rust-lang/crates.io-index"
389+
checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5"
390+
dependencies = [
391+
"generic-array",
392+
]
393+
335394
[[package]]
336395
name = "instant"
337396
version = "0.1.12"
@@ -469,6 +528,7 @@ dependencies = [
469528
name = "miri"
470529
version = "0.1.0"
471530
dependencies = [
531+
"aes",
472532
"colored",
473533
"ctrlc",
474534
"env_logger",
@@ -909,6 +969,12 @@ dependencies = [
909969
"tracing-core",
910970
]
911971

972+
[[package]]
973+
name = "typenum"
974+
version = "1.17.0"
975+
source = "registry+https://github.com/rust-lang/crates.io-index"
976+
checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"
977+
912978
[[package]]
913979
name = "ui_test"
914980
version = "0.21.2"
@@ -954,6 +1020,12 @@ version = "0.1.0"
9541020
source = "registry+https://github.com/rust-lang/crates.io-index"
9551021
checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d"
9561022

1023+
[[package]]
1024+
name = "version_check"
1025+
version = "0.9.4"
1026+
source = "registry+https://github.com/rust-lang/crates.io-index"
1027+
checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
1028+
9571029
[[package]]
9581030
name = "wasi"
9591031
version = "0.11.0+wasi-snapshot-preview1"

src/tools/miri/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ env_logger = "0.10"
2323
log = "0.4"
2424
rand = "0.8"
2525
smallvec = "1.7"
26+
aes = { version = "0.8.3", features = ["hazmat"] }
2627

2728
measureme = "10.0.0"
2829
ctrlc = "3.2.5"

src/tools/miri/src/shims/x86/aesni.rs

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
use rustc_middle::ty::layout::LayoutOf as _;
2+
use rustc_middle::ty::Ty;
3+
use rustc_span::Symbol;
4+
use rustc_target::spec::abi::Abi;
5+
6+
use crate::*;
7+
use shims::foreign_items::EmulateForeignItemResult;
8+
9+
impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
10+
pub(super) trait EvalContextExt<'mir, 'tcx: 'mir>:
11+
crate::MiriInterpCxExt<'mir, 'tcx>
12+
{
13+
fn emulate_x86_aesni_intrinsic(
14+
&mut self,
15+
link_name: Symbol,
16+
abi: Abi,
17+
args: &[OpTy<'tcx, Provenance>],
18+
dest: &PlaceTy<'tcx, Provenance>,
19+
) -> InterpResult<'tcx, EmulateForeignItemResult> {
20+
let this = self.eval_context_mut();
21+
// Prefix should have already been checked.
22+
let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.aesni.").unwrap();
23+
24+
match unprefixed_name {
25+
// Used to implement the _mm_aesdec_si128, _mm256_aesdec_epi128
26+
// and _mm512_aesdec_epi128 functions.
27+
// Performs one round of an AES decryption on each 128-bit word of
28+
// `state` with the corresponding 128-bit key of `key`.
29+
// https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesdec_si128
30+
"aesdec" | "aesdec.256" | "aesdec.512" => {
31+
let [state, key] =
32+
this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
33+
34+
aes_round(this, state, key, dest, |state, key| {
35+
let key = aes::Block::from(key.to_le_bytes());
36+
let mut state = aes::Block::from(state.to_le_bytes());
37+
// `aes::hazmat::equiv_inv_cipher_round` documentation states that
38+
// it performs the same operation as the x86 aesdec instruction.
39+
aes::hazmat::equiv_inv_cipher_round(&mut state, &key);
40+
u128::from_le_bytes(state.into())
41+
})?;
42+
}
43+
// Used to implement the _mm_aesdeclast_si128, _mm256_aesdeclast_epi128
44+
// and _mm512_aesdeclast_epi128 functions.
45+
// Performs last round of an AES decryption on each 128-bit word of
46+
// `state` with the corresponding 128-bit key of `key`.
47+
// https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesdeclast_si128
48+
"aesdeclast" | "aesdeclast.256" | "aesdeclast.512" => {
49+
let [state, key] =
50+
this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
51+
52+
aes_round(this, state, key, dest, |state, key| {
53+
let mut state = aes::Block::from(state.to_le_bytes());
54+
// `aes::hazmat::equiv_inv_cipher_round` does the following operations:
55+
// state = InvShiftRows(state)
56+
// state = InvSubBytes(state)
57+
// state = InvMixColumns(state)
58+
// state = state ^ key
59+
// But we need to skip the InvMixColumns.
60+
// First, use a zeroed key to skip the XOR.
61+
aes::hazmat::equiv_inv_cipher_round(&mut state, &aes::Block::from([0; 16]));
62+
// Then, undo the InvMixColumns with MixColumns.
63+
aes::hazmat::mix_columns(&mut state);
64+
// Finally, do the XOR.
65+
u128::from_le_bytes(state.into()) ^ key
66+
})?;
67+
}
68+
// Used to implement the _mm_aesenc_si128, _mm256_aesenc_epi128
69+
// and _mm512_aesenc_epi128 functions.
70+
// Performs one round of an AES encryption on each 128-bit word of
71+
// `state` with the corresponding 128-bit key of `key`.
72+
// https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesenc_si128
73+
"aesenc" | "aesenc.256" | "aesenc.512" => {
74+
let [state, key] =
75+
this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
76+
77+
aes_round(this, state, key, dest, |state, key| {
78+
let key = aes::Block::from(key.to_le_bytes());
79+
let mut state = aes::Block::from(state.to_le_bytes());
80+
// `aes::hazmat::cipher_round` documentation states that
81+
// it performs the same operation as the x86 aesenc instruction.
82+
aes::hazmat::cipher_round(&mut state, &key);
83+
u128::from_le_bytes(state.into())
84+
})?;
85+
}
86+
// Used to implement the _mm_aesenclast_si128, _mm256_aesenclast_epi128
87+
// and _mm512_aesenclast_epi128 functions.
88+
// Performs last round of an AES encryption on each 128-bit word of
89+
// `state` with the corresponding 128-bit key of `key`.
90+
// https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesenclast_si128
91+
"aesenclast" | "aesenclast.256" | "aesenclast.512" => {
92+
let [state, key] =
93+
this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
94+
95+
aes_round(this, state, key, dest, |state, key| {
96+
let mut state = aes::Block::from(state.to_le_bytes());
97+
// `aes::hazmat::cipher_round` does the following operations:
98+
// state = ShiftRows(state)
99+
// state = SubBytes(state)
100+
// state = MixColumns(state)
101+
// state = state ^ key
102+
// But we need to skip the MixColumns.
103+
// First, use a zeroed key to skip the XOR.
104+
aes::hazmat::cipher_round(&mut state, &aes::Block::from([0; 16]));
105+
// Then, undo the MixColumns with InvMixColumns.
106+
aes::hazmat::inv_mix_columns(&mut state);
107+
// Finally, do the XOR.
108+
u128::from_le_bytes(state.into()) ^ key
109+
})?;
110+
}
111+
// Used to implement the _mm_aesimc_si128 function.
112+
// Performs the AES InvMixColumns operation on `op`
113+
"aesimc" => {
114+
let [op] = this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
115+
116+
// Transmute to `u128`
117+
let op = op.transmute(this.machine.layouts.u128, this)?;
118+
let dest = dest.transmute(this.machine.layouts.u128, this)?;
119+
120+
let state = this.read_scalar(&op)?.to_u128()?;
121+
let mut state = aes::Block::from(state.to_le_bytes());
122+
aes::hazmat::inv_mix_columns(&mut state);
123+
124+
this.write_scalar(Scalar::from_u128(u128::from_le_bytes(state.into())), &dest)?;
125+
}
126+
// TODO: Implement the `llvm.x86.aesni.aeskeygenassist` when possible
127+
// with an external crate.
128+
_ => return Ok(EmulateForeignItemResult::NotSupported),
129+
}
130+
Ok(EmulateForeignItemResult::NeedsJumping)
131+
}
132+
}
133+
134+
// Performs an AES round (given by `f`) on each 128-bit word of
135+
// `state` with the corresponding 128-bit key of `key`.
136+
fn aes_round<'tcx>(
137+
this: &mut crate::MiriInterpCx<'_, 'tcx>,
138+
state: &OpTy<'tcx, Provenance>,
139+
key: &OpTy<'tcx, Provenance>,
140+
dest: &PlaceTy<'tcx, Provenance>,
141+
f: impl Fn(u128, u128) -> u128,
142+
) -> InterpResult<'tcx, ()> {
143+
assert_eq!(dest.layout.size, state.layout.size);
144+
assert_eq!(dest.layout.size, key.layout.size);
145+
146+
// Transmute arguments to arrays of `u128`.
147+
assert_eq!(dest.layout.size.bytes() % 16, 0);
148+
let len = dest.layout.size.bytes() / 16;
149+
150+
let u128_array_layout =
151+
this.layout_of(Ty::new_array(this.tcx.tcx, this.tcx.types.u128, len))?;
152+
153+
let state = state.transmute(u128_array_layout, this)?;
154+
let key = key.transmute(u128_array_layout, this)?;
155+
let dest = dest.transmute(u128_array_layout, this)?;
156+
157+
for i in 0..len {
158+
let state = this.read_scalar(&this.project_index(&state, i)?)?.to_u128()?;
159+
let key = this.read_scalar(&this.project_index(&key, i)?)?.to_u128()?;
160+
let dest = this.project_index(&dest, i)?;
161+
162+
let res = f(state, key);
163+
164+
this.write_scalar(Scalar::from_u128(res), &dest)?;
165+
}
166+
167+
Ok(())
168+
}

src/tools/miri/src/shims/x86/mod.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use crate::*;
77
use helpers::bool_to_simd_element;
88
use shims::foreign_items::EmulateForeignItemResult;
99

10+
mod aesni;
1011
mod sse;
1112
mod sse2;
1213
mod sse3;
@@ -100,6 +101,12 @@ pub(super) trait EvalContextExt<'mir, 'tcx: 'mir>:
100101
this, link_name, abi, args, dest,
101102
);
102103
}
104+
name if name.starts_with("aesni.") => {
105+
return aesni::EvalContextExt::emulate_x86_aesni_intrinsic(
106+
this, link_name, abi, args, dest,
107+
);
108+
}
109+
103110
_ => return Ok(EmulateForeignItemResult::NotSupported),
104111
}
105112
Ok(EmulateForeignItemResult::NeedsJumping)

0 commit comments

Comments
 (0)