aes 0.9.0-pre.2 uses ParBlocksSize of 9 but has type alias for Block8

[robinhundt]

I recently used the current pre-release version of the aes crate and stumbled on a performance pitfall. In previous versions, the AES-NI implementation processed 8 blocks in parallel. This was changed in #442 to 9 blocks (see this comment). I chunked an iterator into 8 block chunks but got no speedup due to this.

Unfortunately, aes still exposes the Block8 type alias. Personally, I thought the reason for the type alias was it's usefulness when wanting to make use of parallel processing. Admittedly, this is not described by the Block8 documentation but supported by the hazmat API.

So my first suggestion would be to expand the docs of Block8 to clarify that its existence is not tied to the number of parallel rounds (which depend on the backend anyway). If this type is useful only in the context of the hazmat API, maybe it could even be moved there before the release of 0.9.

While I'm on the topic of ParBlockSize: Is there currently any way to get the ParBlocksSize for a Backend? If not, should I create an issue for this?

[newpavlov]

The hazmat functions still use 8 blocks, i.e. they are not tied to the cipher backend implementations. Providing parallel block size variability would make the functions API significantly more complex and its hard to find an optimal number for them, since it can depend on number of rounds used by a hazmat API user. So I decided to just leave the 8 block API intact.

Is there currently any way to get the ParBlocksSize for a Backend?

Yes, you should be able to use B::ParBlocksSize where B is backend passed to your "closure" implementation.

[robinhundt]

The hazmat functions still use 8 blocks, i.e. they are not tied to the cipher backend implementations.

What is your opinion on moving Block8 into the hazmat module to reduce confusion? (Not sure about the amount of churn this would entail)

Yes, you should be able to use B::ParBlocksSize where B is backend passed to your "closure" implementation.

Thanks for the tip! I noticed the encrypt_with_backend function, but honestly, all those interconnected traits intimidated me a little 😅
I've come up with the following solution:

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
// https://github.com/RustCrypto/block-ciphers/blob/4da9b802de52a3326fdc74d559caddd57042fed2/aes/src/ni.rs#L43
pub const AES_PAR_BLOCKS: usize = 9;
#[cfg(target_arch = "aarch64")]
// https://github.com/RustCrypto/block-ciphers/blob/4da9b802de52a3326fdc74d559caddd57042fed2/aes/src/armv8.rs#L32
pub const AES_PAR_BLOCKS: usize = 21;
#[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
// TODO what should the fallback be?
pub const AES_PAR_BLOCKS: usize = 4;

#[cfg(test)]
mod tests {
    use std::mem;

    use aes::{
        cipher::{
            BlockCipherEncClosure, BlockCipherEncrypt, BlockSizeUser, KeyInit, ParBlocksSizeUser,
        },
        Aes128,
    };
    use hybrid_array::ArraySize;

    use crate::AES_PAR_BLOCKS;

    #[cfg(target_feature = "aes")]
    #[test]
    fn aes_par_block_size() {
        struct GetParBlockSize;
        impl BlockSizeUser for GetParBlockSize {
            type BlockSize = aes::cipher::array::sizes::U16;
        }
        impl BlockCipherEncClosure for GetParBlockSize {
            fn call<B: aes::cipher::BlockCipherEncBackend<BlockSize = Self::BlockSize>>(
                self,
                _backend: &B,
            ) {
                assert_eq!(
                    AES_PAR_BLOCKS,
                    // size_of ArrayType<u8> is equal to its length
                    mem::size_of::<
                        <<B as ParBlocksSizeUser>::ParBlocksSize as ArraySize>::ArrayType<u8>,
                    >()
                );
            }
        }
        let aes = Aes128::new(&Default::default());
        aes.encrypt_with_backend(GetParBlockSize);
    }
}

This way I can use the ParBlocksSize as a constant, but don't get surprised when it changes.

[newpavlov]

What is your opinion on moving Block8 into the hazmat module to reduce confusion?

It may be a worthwhile change. @tarcieri WDYT?

Thanks for the tip! I noticed the encrypt_with_backend function, but honestly, all those interconnected traits intimidated me a little

Yeah, the API is somewhat hard to understand at the first glance. You can imagine writing code like this:

let cipher = Aes128::new(&key);
cipher.encrypt_with_backend(|cipher_backend| {
    // This closure captures some external state, e.g. blocks which we want to encrypt.
    // `cipher_backend` provides `encrypt_block` and `encrypt_par_blocks` methods
    // and has a concrete "parallel block size" accessible via `cipher_backend::ParBlockSize`
});

Unfortunately, current version of Rust does not allow definition of methods which accept rank-2 closures, so we have to effectively implement closures manually instead of relying on the convenient |args| { .. } syntax. The goal here is to generate a copy of the closure code for each available backend.

I've come up with the following solution

I don't think it's a good solution. In future we may change parallel block size or introduce new backends (e.g. based on AVX-512) and your code may break. Even worse, it may break depending on target features which are available on your user's CPU.

[robinhundt]

I'm not relying on the ParBlocksSize for safety or correctness, only to squeeze out the maximum performance. In the worst case, a wrong constant should only result in lower performance because the parallel methods are not used due to a wrong chunk size. This is also not used for production code.

So I want to use the most likely ParSizeBlocks for my current backend. That's why I also didn't cfg the constant on the aes feature, as my crate might not be compiled with aes, but if the system has AES-NI, the runtime detection of the aes crate should use these intrinsics when available, right?
I guess I could add a method that uses the above trick to get the ParBlocksSize of the current backend at runtime for chunking, but that doesn't help me in the cases where I need a constant for e.g. array lengths.

I'm totally fine with my test breaking when you change the ParBlocksSize or introduce something new. But if you have a different suggestion on how I could (semi) reliably get the ParBlocksSize, I'd be happy to hear it.

[newpavlov]

Well, for experimental non-production code I guess such approach could be fine, but I would say it smells a bit.

If you target a known hardware, it may make sense to use AES-NI directly in your code since resulting code is quite small and relatively straightforward. I am not sure we want to expose it in aes public API because of backward compatibility concerns (maybe behind an "unstable" crate feature?).

I guess I could add a method that uses the above trick to get the ParBlocksSize of the current backend at runtime for chunking, but that doesn't help me in the cases where I need a constant for e.g. array lengths.

You can get usize from ParBlocksSize by using the USIZE associated constant on the typenum::Unsigned trait, i.e. B::ParBlocksSize::USIZE. Granted, you will not be able to use this constant for creating built-in arrays, but you should be able to use it to create hybrid_array::Array (this only applies to aes v0.9, on v0.8 we use generic-array which leads to somewhat stricter restrictions).

[robinhundt]

If you target a known hardware, it may make sense to use AES-NI directly in your code since resulting code is quite small and relatively straightforward. I am not sure we want to expose it in aes public API because of backward compatibility concerns (maybe behind an "unstable" crate feature?).

Thanks for the tip! While I'm mostly targeting hardware with AES-NI, I would like my code to be portable across most targets. But I think I have a good use case where I can use the hazmat API for more flexibility instead of the normal one.

You can get usize from ParBlocksSize by using the USIZE associated constant on the typenum::Unsigned trait, i.e. B::ParBlocksSize::USIZE. Granted, you will not be able to use this constant for creating built-in arrays, but you should be able to use it to create hybrid_array::Array (this only applies to aes v0.9, on v0.8 we use generic-array which leads to somewhat stricter restrictions).

Thanks, I tried that first but got the trait casts wrong. I looked again, and <<B as ParBlocksSizeUser>::ParBlocksSize as Unsigned>::USIZE works :)