Skip to main content

aes/backends/
x86_vaes512.rs

1use super::x86_aes::RoundKeys;
2use cipher::{
3    Block, BlockCipherDecBackend, BlockCipherDecClosure, BlockCipherEncBackend,
4    BlockCipherEncClosure, BlockSizeUser, ParBlocks, ParBlocksSizeUser,
5    consts::{U16, U64},
6    inout::InOut,
7};
8
9mod encdec;
10
11pub(crate) type Aes128<'a> = Aes<'a, 11>;
12pub(crate) type Aes192<'a> = Aes<'a, 13>;
13pub(crate) type Aes256<'a> = Aes<'a, 15>;
14
15#[derive(Clone, Copy)]
16pub(crate) struct Aes<'a, const RK: usize> {
17    rk: &'a RoundKeys<RK>,
18    rk4: encdec::RoundKeys4<RK>,
19}
20
21impl<'a, const RK: usize> Aes<'a, RK> {
22    #[inline]
23    #[target_feature(enable = "avx512f,vaes")]
24    pub(crate) fn encrypt(rk: &'a RoundKeys<RK>, f: impl BlockCipherEncClosure<BlockSize = U16>) {
25        let rk4 = encdec::broadcast_keys(rk);
26        let backend = Self { rk, rk4 };
27        f.call(&backend)
28    }
29
30    #[inline]
31    #[target_feature(enable = "avx512f,vaes")]
32    pub(crate) fn decrypt(rk: &'a RoundKeys<RK>, f: impl BlockCipherDecClosure<BlockSize = U16>) {
33        let rk4 = encdec::broadcast_keys(rk);
34        let backend = Self { rk, rk4 };
35        f.call(&backend)
36    }
37}
38
39impl<const RK: usize> BlockSizeUser for Aes<'_, RK> {
40    type BlockSize = U16;
41}
42
43// Block sizes are chosen based on AVX-512's 32 ZMM registers.
44//
45// Round keys use 11, 13, 15 registers for AES-128/192/256 respectively.
46// It results in 21, 19, 17 registers available for blocks with each register containing 4 blocks.
47// We use the closest power-of-two value.
48// TODO: bench bigger cipher-specific values
49impl<const RK: usize> ParBlocksSizeUser for Aes<'_, RK> {
50    type ParBlocksSize = U64;
51}
52
53impl<const RK: usize> BlockCipherEncBackend for Aes<'_, RK> {
54    #[inline(always)]
55    fn encrypt_block(&self, block: InOut<'_, '_, Block<Self>>) {
56        // SAFETY: this trait impl is used only by the `Self::encrypt` method marked with
57        // `#[target_feature(enable = "vaavx512f,vaeses")]`
58        unsafe { super::x86_aes::encrypt(self.rk, block) };
59    }
60
61    #[inline(always)]
62    fn encrypt_par_blocks(&self, blocks: InOut<'_, '_, ParBlocks<Self>>) {
63        // SAFETY: this trait impl is used only by the `Self::encrypt` method marked with
64        // `#[target_feature(enable = "avx512f,vaes")]`
65        unsafe { encdec::batch_encrypt(&self.rk4, blocks) };
66    }
67}
68
69impl<const RK: usize> BlockCipherDecBackend for Aes<'_, RK> {
70    #[inline(always)]
71    fn decrypt_block(&self, block: InOut<'_, '_, Block<Self>>) {
72        // SAFETY: this trait impl is used only by the `Self::decrypt` method marked with
73        // `#[target_feature(enable = "avx512f,vaes")]`
74        unsafe { super::x86_aes::decrypt(self.rk, block) };
75    }
76
77    #[inline(always)]
78    fn decrypt_par_blocks(&self, blocks: InOut<'_, '_, ParBlocks<Self>>) {
79        // SAFETY: this trait impl is used only by the `Self::decrypt` method marked with
80        // `#[target_feature(enable = "avx512f,vaes")]`
81        unsafe { encdec::batch_decrypt(&self.rk4, blocks) };
82    }
83}