aes: SIMD-acceleration for WASM

[sinui0]

WASM targets currently use the software fixslice implementation which processes up to 4 blocks at a time (fixslice64). WebAssembly has had the simd128 proposal stabilized in the spec for a while now, and the corresponding intrinsics ([core::arch::wasm32::*]) have been stable in Rust since 1.33. Every major WebAssembly runtime ships simd128 support today.

We use AES extensively in our projects which target running in browsers, the throughput of AES in WASM is a major performance lever for us. Translating the existing fixslice64 implementation (AI assisted, not sure what your policy is on that) to use 128 bit lanes yields the expected >2x throughput increase in V8 (chromium), and >3x in SpiderMonkey (firefox).

(AMD Ryzen 7 5800X)

Unit: MB/s	wasmtime soft	wasmtime simd128	wasmtime relative	v8 soft	v8 simd128	v8 relative	sm soft	sm simd128	sm relative
AES-128 enc	186	184	0.99×	258	604	2.34×	213	657	3.08×
AES-128 dec	185	122	0.66×	235	508	2.16×	182	598	3.29×
AES-192 enc	159	150	0.94×	224	522	2.33×	185	588	3.18×
AES-192 dec	159	100	0.63×	202	486	2.41×	153	577	3.77×
AES-256 enc	138	140	1.01×	192	455	2.37×	159	491	3.09×
AES-256 dec	135	89	0.66×	171	424	2.48×	134	453	3.38×

As you can see, this causes a regression in Wasmtime decrypt performance due to bad JIT by Cranelift. I haven't been able to mitigate that without incurring significant throughput loss in the browser runtimes. I'm not sure whether the right approach is to eat the regression by default, or to put the SIMD behind a flag.

Implementation

Can reproduce the benchmarks using wasm-harness

Would love to upstream this if possible. If so, I will open a PR.

[newpavlov]

Have you tried to just bump the FixsliceBlocks constant and let the compiler handle the rest?

[sinui0]

I'm not sure if I fully understand your suggestion, but I tried bumping it up to 8 and interleaving two 4 block batches together to see if LLVM would auto-vectorize it. No success, and I would be very impressed if it were able to do so.

Note: just bumping that constant will cause the implementation to only encrypt the first 4 blocks and leave the rest as plaintext

[tarcieri]

It'd be really nice if auto-vectorization would work here. I'd suggest trying to break the problem down in something like godbolt and seeing if maybe you could get it to work with e.g. the fixslicing schedule before trying to get it working with the full AES.

The alternative of having an alternate fixslicing backend that uses an explicit v128 instead of u64 seems like it will work although there is already an annoying amount of duplication between the 32-bit and 64-bit backends and this sounds like it will further compound that.

I would be curious if the implementation could be made generic somehow so you could at least reuse large portions of the same code for u64 and v128.

[sinui0]

I agree that achieving auto-vectorization would be preferable. I gave it a shot, and in my experiments LLVM will autovectorize much of AES at 128-bit width, except the S-Box circuit which seems to need an adjustment to LLVM thresholds otherwise it gives up being too high depth.

My current thinking is that the fixslice algorithm can be parameterized over the lane width and the vast majority of it can be rewritten to be generic over a Lane abstraction so that 32/64/128 width lane impls differ only in a small number of places. This would address your duplication/maintenance concern and likely net reduce LOC even with a 128 bit variant added.

With that, I do think that a WASM platform specialization would be needed to ensure the S-Box is vectorized. With the above refactor this should only need a few hundred LOC.

Is this something you would accept? I'm cognizant that the AES impl has been audited, so I would understand an aversion to touching it. For us, the performance of the WASM build is critical so we're motivated to get this 2-3x for it.

[tarcieri]

Is this something you would accept?

It sounds like a potentially interesting direction, but I'd need to see the actual implementation to know