Add initial vectorized chacha20 implementation for 2-3x speedup #34083

Why separate this into an .ipp file if it's only included in a single .cpp file?

See branch here: https://github.com/theuni/bitcoin/commits/chacha20-vectorized/

I decided to exclude the impls which require runtime detection from this PR, as I think how that should be done is a separate conversation.

To answer your question more specifically: some impls may require different compilation flags (-mavx2/-mavx512vl), which have to be in their own compilation units.

Using #if !defined for these seems old school. Why not static constexpr bools, and if constexpr (..)?

Writing:

#if defined(__x86_64__) || defined(__amd64__)
#  define CHACHA20_VEC_DISABLE_STATES_16 true
#  define CHACHA20_VEC_DISABLE_STATES_8  true
#  define CHACHA20_VEC_DISABLE_STATES_6  true
#  define CHACHA20_VEC_DISABLE_STATES_4  false
#  define CHACHA20_VEC_DISABLE_STATES_2  false
#elif defined(__ARM_NEON)
...

static constexpr bool CHACHA20_VEC_ALL_MULTI_STATES_DISABLED =
    CHACHA20_VEC_DISABLE_STATES_16 &&
    CHACHA20_VEC_DISABLE_STATES_8 &&
    CHACHA20_VEC_DISABLE_STATES_6 &&
    CHACHA20_VEC_DISABLE_STATES_4 &&
    CHACHA20_VEC_DISABLE_STATES_2;

 void chacha20_crypt_vectorized(std::span<const std::byte>& in_bytes, std::span<std::byte>& out_bytes, const std::array<uint32_t, 12>& input) noexcept
 {
    if constexpr (CHACHA20_VEC_ALL_MULTI_STATES_DISABLED) return;
...
    if constexpr (!CHACHA20_VEC_DISABLE_STATES_16) {
         while(in_bytes.size() >= CHACHA20_VEC_BLOCKLEN * 16) {
...

seems to work right, and also seems like it avoids needing to put all the potentially unused code in a #if block.

Here's a branch that replaces the #defines and abstracts the recursive template stuff a bit more for your perusal https://github.com/ajtowns/bitcoin/commits/202512-pr34083-templates/

+1 for if constexpr, should simplify the code a lot (especially if we migrate from recursive templates as well)

if constexpr certainly makes more sense where possible. I'll have a look, thanks!

static_assert(sizeof(vec256) == 32) ? (Or sizeof(vec256) == 32 || ALL_MULTI_STATES_DISABLED)

I haven't used this before but the internets tells me we could use the attribute syntax here, something like:

template <typename T, size_t N>
using VectorType [[gnu::vector_size(sizeof(T) * N)]] = T;

using vec256 = VectorType<uint32_t, 8>;

Because they're compiler-specific, the code makes no assumptions about the size/structure/alignment of vec256. The only accesses are via operator[]. So afaik, there's no need to check for this.

Sorry, I meant just so you get an error if you're using a compiler ignores the __attribute__ entirely but still passes the #if in _base.cpp.

if constexpr (ITER + 1 < I) (we have a space after if constexpr elsewhere)

My understanding is that you're inverting the logic here: instead of

#define QUARTERROUND(a,b,c,d) \
    X(a,b,d,16) X(c,d,b,12) X(a,d,b,8) X(c,d,b,7)

This does the X(a,b,d,16) all four times via simd, then X(c,d,b,12) all four times, etc. And the simd part relies on the data for those four steps being exactly +4 units apart, which is why the shuffling and unshuffling is necessary for the second round. (Okay, not four times but eight times because each vec256 is two blocks)

Could use a little more explanation in the comment I think? But makes sense to me.

Makes sense to me too, but I'm not sure why doing it this way is preferable. Intuitively I would have expected that the vectors hold the same word spread over multiple blocks. But if I understand your approach here, we have multiple words over two blocks. Is it just easier to express the quarter rounds in this way?

Overflow happens every 274GB I guess, which presumably isn't worth putting much effort in to optimising around.

This is actually doing the opposite.. it's keeping the vectorized impl from having to worry about this case. In the current code, we have:

++j12;
if (!j12) ++j13;
...
input[8] = j12;
input[9] = j13;

So effectively input[8] and input[9] are treated as a single uint64_t. It's not possible to express "cast to uint64_t elements and increment" or "increment and overflow over there" with the vector extensions, so it turned out to be easier to just forbid the overflow cases from being vectorized at all.

Yeah, I was thinking about whether the code should just do the non-vectorized stuff to get past the overflow then immediately go back to vectorizing, rather than waiting for the next call. I think what you've got makes sense.

can we assume that all big endian systems have __builtin_bswap32 available? https://github.com/bitcoin/bitcoin/blob/432b18ca8d0654318a8d882b28b20af2cb2d2e5d/src/compat/byteswap.h#L14-L16 indicates we could use our existing helpers here instead:

ALWAYS_INLINE void vec_byteswap(vec256& vec)
{
    if constexpr (std::endian::native == std::endian::big) {
        for (size_t i = 0; i < 8; ++i) {
            vec[i] = internal_bswap_32(vec[i]);
        }
    }
}

A small, fixed-size loop like this should be unrolled by every compiler, it's what we did in https://github.com/bitcoin/bitcoin/pull/31144/files#diff-f26d4597a7f5a5d4aa30053032d49427b402ef4fd7a1b3194cb75b2551d58ca4R48 as well.

(nit: could you please reformat the patch, it differs slightly from how clang-format is set for new code)

Thanks, yes, I meant to change that to internal_bswap_32 before pushing. Will do.

It seems to me the existing standard test vectors are too short (mostly < 128 bytes) to trigger the vectorized optimization - can we extend the tests to runa and compare the new specializations (could show as skipped when the given architecture isn't available locally)?

Yep, will do.

Forcing vectorized operations on an emulated big-endian system (otherwise it falls back to the scalar implementation due to the safety checks):

diff --git a/src/crypto/chacha20_vec_base.cpp b/src/crypto/chacha20_vec_base.cpp
index 9fda9452a1..a2aa1e5552 100644
--- a/src/crypto/chacha20_vec_base.cpp
+++ b/src/crypto/chacha20_vec_base.cpp
@@ -20,7 +20,7 @@
 #  define CHACHA20_VEC_DISABLE_STATES_8
 #  define CHACHA20_VEC_DISABLE_STATES_6
 #  define CHACHA20_VEC_DISABLE_STATES_4
-#  define CHACHA20_VEC_DISABLE_STATES_2
+//#  define CHACHA20_VEC_DISABLE_STATES_2
 #endif

 #include <crypto/chacha20_vec.ipp>

and running the crypto_tests:

brew install podman pigz qemu
podman machine init
podman machine start

podman run --platform linux/s390x -it --rm ubuntu:latest /bin/bash -c \
  'apt-get update && \
   DEBIAN_FRONTEND=noninteractive apt-get install -y \
   git build-essential cmake ccache pkg-config \
   libevent-dev libboost-dev libssl-dev libsqlite3-dev python3 && \
   git clone https://github.com/bitcoin/bitcoin.git && cd bitcoin && \
   git fetch origin pull/34083/head:chacha20-vec && git checkout chacha20-vec && \
   sed -i "s/#  define CHACHA20_VEC_DISABLE_STATES_2/\/\/#  define CHACHA20_VEC_DISABLE_STATES_2/" src/crypto/chacha20_vec_base.cpp && \
   cmake -B build -DBUILD_BENCH=OFF -DBUILD_GUI=OFF -DBUILD_DAEMON=OFF -DBUILD_TX=OFF -DENABLE_IPC=OFF -DENABLE_WALLET=OFF -DENABLE_ZMQ=OFF -DENABLE_UPNP=OFF -DENABLE_NATPMP=OFF && \
   cmake --build build --target test_bitcoin -j1 && \ ./build/bin/test_bitcoin --run_test=crypto_tests'

We're getting a lot of failures:

Running 17 test cases...
./test/crypto_tests.cpp(155): error: in "crypto_tests/chacha20_testvector": check hexout == HexStr(outres) has failed [a3fbf07df3fa2fde4f376ca23e82737041605d9f4f4f57bd8cff2c1d4b7955ec2a97948bd3722915c8f3d337f7d370050e9e96d647b7c39f56e031ca5eb6250d4042e02785ececfa4b4bb5e8ead0440e20b6e8db09d881a7c6132f420e52795042bdfa7773d8a9051447b3291ce1411c680465552aa6c405b7764d5e87bea85ad00f8449ed8f72d0d662ab052691ca66424bc86d2df80ea41f43abf937d3259dc4b2d0dfb48a6c9139ddd7f76966e928e635553ba76c5c879d7b35d49eb2e62b0871cdac638939e25e8a1e0ef9d5280fa8ca328b351c3c765989cbcf3daa8b6ccc3aaf9f3979c92b3720fc88dc95ed84a1be059c6499b9fda236e7e818b04b0bc39c1e876b193bfe5569753f88128cc08aaa9b63d1a16f80ef2554d7189c411f5869ca52c5b83fa36ff216b9c1d30062bebcfd2dc5bce0911934fda79a86f6e698ced759c3ff9b6477338f3da4f9cd8514ea9982ccafb341b2384dd902f3d1ab7ac61dd29c6f21ba5b862f3730e37cfdc4fd806c22f221 != c2ece71ceded38c04f376ca225cc3d6b463409986f263e9db1994a204b6844ec6e9695cfc52b7003e3b6961ceac96a18138981cb4ee5919649b263d542b82b114a57f52d8ba2a2f4540af4f7f49f0b1072a7f9891b97ceb5c61c20420143685f169add2373c4b505122eda2f5df3535d555637680dc5a722b83209519ae7f147c11a9158f48479c9926ea7412ac69d6a584ac97768e412e1514fa7b737ce389dc4bbd9df9cc422b95ccfc5926171fe20e9284834a77646878d7a34c485b3e7300c3589a8448b3bc53b980c6bced42938afc1398c2f1a3a6c2887c5be68a18039cb2fba983271c1202a73af95c79ae29faafb40973694b4afa523f2ef13b84300c39c1e876b193bfe5569753f88128cc08aaa9b63d1a16f80ef2554d7189c411f5869ca52c5b83fa36ff216b9c1d30062bebcfd2dc5bce0911934fda79a86f6e698ced759c3ff9b6477338f3da4f9cd8514ea9982ccafb341b2384dd902f3d1ab7ac61dd29c6f21ba5b862f3730e37cfdc4fd806c22f221]
...
./test/crypto_tests.cpp(185): error: in "crypto_tests/chacha20_testvector": check hexout == HexStr(outres) has failed [a3fbf07df3fa2fde4f376ca23e82737041605d9f4f4f57bd8cff2c1d4b7955ec2a97948bd3722915c8f3d337f7d370050e9e96d647b7c39f56e031ca5eb6250d4042e02785ececfa4b4bb5e8ead0440e20b6e8db09d881a7c6132f420e52795042bdfa7773d8a9051447b3291ce1411c680465552aa6c405b7764d5e87bea85ad00f8449ed8f72d0d662ab052691ca66424bc86d2df80ea41f43abf937d3259dc4b2d0dfb48a6c9139ddd7f76966e928e635553ba76c5c879d7b35d49eb2e62b0871cdac638939e25e8a1e0ef9d5280fa8ca328b351c3c765989cbcf3daa8b6ccc3aaf9f3979c92b3720fc88dc95ed84a1be059c6499b9fda236e7e818b04b0bc39c1e876b193bfe5569753f88128cc08aaa9b63d1a16f80ef2554d7189c411f5869ca52c5b83fa36ff216b9c1d30062bebcfd2dc5bce0911934fda79a86f6e698ced759c3ff9b6477338f3da4f9cd8514ea9982ccafb341b2384dd902f3d1ab7ac61dd29c6f21ba5b862f3730e37cfdc4fd806c22f221 != a3fbf07df3fa2fde4f376ca23e82737041605d9f4f4f57bd8cff2c1d4b7955ec2a97948bd3722915c8f3d337f7d370050e9e96d647b7c39f56e031ca5eb6250d4a57f52d8ba2a2f4540af4f7f49f0b1072a7f9891b97ceb5c61c20420143685f169add2373c4b505122eda2f5df3535d555637680dc5a722b83209519ae7f147c11a9158f48479c9926ea7412ac69d6a584ac97768e412e1514fa7b737ce389dc4bbd9df9cc422b95ccfc5926171fe20e9284834a77646878d7a34c485b3e7300871cdac638939e25e8a1e0ef9d5280fa8ca328b351c3c765989cbcf3daa8b6ccc3aaf9f3979c92b3720fc88dc95ed84a1be059c6499b9fda236e7e818b04b0bc39c1e876b193bfe5569753f88128cc08aaa9b63d1a16f80ef2554d7189c411f5869ca52c5b83fa36ff216b9c1d30062bebcfd2dc5bce0911934fda79a86f6e698ced759c3ff9b6477338f3da4f9cd8514ea9982ccafb341b2384dd902f3d1ab7ac61dd29c6f21ba5b862f3730e37cfdc4fd806c22f221]
./test/crypto_tests.cpp(272): error: in "crypto_tests/chacha20poly1305_testvectors": check cipher == expected_cipher has failed
...
./test/crypto_tests.cpp(337): error: in "crypto_tests/chacha20poly1305_testvectors": check decipher == plain has failed

*** 37 failures are detected in the test module "Bitcoin Core Test Suite"

This appears to be an endianness issue in vec_read_xor_write: the current implementation swaps the result of the XOR, but on Big Endian systems we must swap the state vec before the XOR.

Changing it to:

diff --git a/src/crypto/chacha20_vec.ipp b/src/crypto/chacha20_vec.ipp
index 46a159ce01..cfc0535a92 100644
--- a/src/crypto/chacha20_vec.ipp
+++ b/src/crypto/chacha20_vec.ipp
@@ -174,8 +174,9 @@ ALWAYS_INLINE void vec_read_xor_write(std::span<const std::byte, 32> in_bytes, s
 {
     std::array<uint32_t, 8> temparr;
     memcpy(temparr.data(), in_bytes.data(), in_bytes.size());
-    vec256 tempvec = vec ^ (vec256){temparr[0], temparr[1], temparr[2], temparr[3], temparr[4], temparr[5], temparr[6], temparr[7]};
+    vec256 tempvec = vec;
     vec_byteswap(tempvec);
+    tempvec ^= (vec256){temparr[0], temparr[1], temparr[2], temparr[3], temparr[4], temparr[5], temparr[6], temparr[7]};
     temparr = {tempvec[0], tempvec[1], tempvec[2], tempvec[3], tempvec[4], tempvec[5], tempvec[6], tempvec[7]};
     memcpy(out_bytes.data(), temparr.data(), out_bytes.size());
 }

Makes it pass for me.

We should find a way to exercise this via CI and benchmarks, maybe similarly to SHA256AutoDetect in https://github.com/bitcoin/bitcoin/blob/bdb8eadcdc193f398ebad83911d3297b5257e721/src/crypto/sha256.cpp#L585-L690 which would enable us running benchmarks and tests selectively.

Thanks for catching this! I forgot to mention in the PR description that big-endian was best-effort and untested. I figured our c-i would catch any obvious bugs. Agree it's not great that it didn't :(

Thanks for the quick fix too :)

I had the same disappointment when implementing the obfuscation optimization. :) @maflcko do we have a big-endian nightly that supports vectorized operations? Would it have caught this?

See #33436, but it was slow despite having the gui and the fuzz tests disabled. I am running it as part of nightly, so any issues will be caught before a release, but I am not sure if catching everything in pull requests is possible.

Maybe it is possible to split the task into two: One for a cross-compile, which should be faster than a "native" compile via qemu. And another to run the tests, similar to the windows-cross tests.

Actually, I ran the CI config locally, but it didn't catch this, as the platform opts out. The CI doesn't run the sed -i "s/# define CHACHA20_VEC_DISABLE_STATES_2/\/\/# define CHACHA20_VEC_DISABLE_STATES_2/" src/crypto/chacha20_vec_base.cpp && \ portion, so it would not have caught this, even if the task was run.

Thanks for checking @maflcko, that's why I asked. The vectorized operations are obviously supported, but for some reason were not triggered for me either. @theuni, is this just an emulation anomaly or we were just too cautious? I understdood that @achow101 has access to real big-endian power9 machine that we might be able to test this on when it's ready.

@l0rinc As the code is written at the moment, all platforms must opt-in to vectorization as opposed to opting out. I'm not sure that's the best approach, but I figured that was a reasonable starting point.

My reasoning for that was: consider non-x86, non-arm+neon platforms. For the most part, I'm assuming they're under-powered. Enabling (for example) 4x blocks/sec for mipsel would probably cause a drastic slowdown. Obviously there are lots of other powerful platforms, but I figured those would be added over time.

So if there's a big-endian architecture (or more ideally, a specific instruction set ala __ARM_NEON) that demonstrates a performance gain from calculating multiple states at once, it should be added here.

Of course, we could go the opposite direction and opt all platforms IN to all states, and instead opt-out the slow ones.

tl;dr: If we want a big-endian platform to be supported, we need to opt one in or provide an override.

Hmmm, it seems to me we're doing heavy recursive templates here to unroll loops. My understanding (and experience with the mentioned Obfuscation PR) is that modern compilers are good at unrolling fixed-bound loops. Can you please try if this also works and results in the same speedup?

template <size_t BITS, size_t I>
ALWAYS_INLINE void arr_add_xor_rot(std::array<vec256, I>& arr0, const std::array<vec256, I>& arr1, std::array<vec256, I>& arr2)
{
    for (size_t i{0}; i < I; ++i) {
        arr0[i] += arr1[i];
        arr2[i] ^= arr0[i];
        vec_rotl<BITS>(arr2[i]);
    }
}

(nit: the size is often N instead of I)

Alternatively, with a compile-time loop:

/** Perform add/xor/rotate for the round function */
template <size_t BITS, size_t I>
ALWAYS_INLINE void arr_add_xor_rot(std::array<vec256, I>& arr0, const std::array<vec256, I>& arr1, std::array<vec256, I>& arr2)
{
    [&]<size_t... ITER>(std::index_sequence<ITER...>) {
        ((
            arr0[ITER] += arr1[ITER],
            arr2[ITER] ^= arr0[ITER],
            vec_rotl<BITS>(arr2[ITER])
        ), ...);
    }(std::make_index_sequence<I>());
}

See updated title description where I touched on this.

I found (with lots of experimentation here) that different compilers use complicated and unpredictable heuristics to decide whether or not to unroll loops. Even if an unroll-able loop is detected, unrolling may be skipped because of the function size (as is the case here because it's huge).

So, I'd like to not take chances on unrolling. That means one of:

• Manual unrolling
• Macro-based (as-in REPEAT10())
• A pragma
• Recursive template based

c++26 introduces template for, which is what we really want.

I'm up for whichever of those is generally preferred, as long as it's explicit rather than implicit.

Is there any way we could help that would make you reconsider? I don't mind running the benchmarks on a few platforms, as long as we can have simpler code. The current template magic is not something I would like to see more of - modern C++20 should be able to handle the situations we have here. I don't mind experimenting with this if you don't want to.

@theuni See my std::make_index_sequence based I demonstrated above. It's template based, but doesn't need recursion, and doesn't rely on compiler unrolling. It simply expands to an expression (a[0] = v, a[1] = v, a[2] = v, ...) e.g.

https://www.agner.org/optimize/optimizing_cpp.pdf contains a few useful examples:

The automatic vectorization works best if the following conditions are satisfied:

1. Use a compiler with good support for automatic vectorization, such as Gnu, Clang, or Intel.

2. Use the latest version of the compiler. The compilers are becoming better and better at vectorization.

3. If the arrays or structures are accessed through pointers or references then tell the compiler explicitly that pointers do not alias, if appropriate, using the restrict or __restrict keyword.

4. Use appropriate compiler options to enable the desired instruction set (/arch:SSE2, /arch:AVX etc. for Windows, -msse2, -mavx512f, etc. for Linux)

5. Use the less restrictive floating point options. For Gnu and Clang compilers, use the options -O2 -fno-trapping-math -fno-math-errno -fno-signed-zeros (-ffast-math works as well, but functions like isnan(x) do not work under -ffast-math).

6. Align arrays and big structures by 16 for SSE2, preferably 32 for AVX and preferably 64 for AVX512.

7. The loop count should preferably be a constant that is divisible by the number of elements in a vector.

8. If arrays are accessed through pointers so that the alignment is not visible in the scope of the function where you want vectorization then follow the advice given above.

9. Minimize the use of branches at the vector element level.

10. Avoid table lookup at the vector element level.

Similarly in https://en.algorithmica.org/hpc/simd/auto-vectorization useful hints:

The other way, specific to SIMD, is the “ignore vector dependencies” pragma. It is a general way to inform the compiler that there are no dependencies between the loop iterations:

#pragma GCC ivdep for (int i = 0; i < n; i++)

Could we use https://github.com/bitcoin/bitcoin/blob/e16c22fe025f82166c7f3f15a37c96bf4a06e4cf/src/attributes.h#L19-L25 instead?

There's a lot of repetition here - could we extract that to an ALWAYS_INLINE lambda and have something like:

    if constexpr (ENABLE_16) process_blocks(16);
    else if constexpr (ENABLE_8)  process_blocks(8);
...

I haven't implemented it locally, maybe it's naive, let me know what you think.

I believe it's possible to use compile-time loops here instead of recursive templates (if runtime loops don't get optimized sufficiently):

/** Store a vector in all array elements */
template <size_t I>
ALWAYS_INLINE void arr_set_vec256(std::array<vec256, I>& arr, const vec256& vec)
{
    [&]<size_t... ITER>(std::index_sequence<ITER...>) {
        ((std::get<ITER>(arr) = vec),...);
    }(std::make_index_sequence<I>());
}

I avoided lambdas as I've historically observed lots of optims being skipped (with clang, at least) when using them. But since you/@ajtowns/@l0rinc have all made the same comment, I'll play around and see if these indeed compile down to nothing as one would hope.

It can be done with a helper function too, but that's not as concise:

template <size_t I, size_t... ITER>
ALWAYS_INLINE void arr_set_vec256_inner(std::array<vec256, I>& arr, const vec256& vec, std::index_sequence<ITER...>)
{
    ((std::get<ITER>(arr) = vec),...);
}

/** Store a vector in all array elements */
template <size_t I>
ALWAYS_INLINE void arr_set_vec256(std::array<vec256, I>& arr, const vec256& vec)
{
    arr_set_vec256_inner(arr, vec, std::make_index_sequence<I>());
}

Writing this as:

    template<size_t... ITER> using iseq = std::index_sequence<ITER>;
    static constexpr ISEQ = std::make_index_sequence<I>();
    
    template<size_t... ITER>
    ALWAYS_INLINE void arr_set_vec256(iseq<ITER>, std::array<vec256, I>& arr, const vec256& vec)
    {
        ((std::get<ITER>(arr) = vec), ...);
    }

   ...
        arr_set_vec256(ISEQ, arr0, num256);

doesn't seem too bad, and avoids lambdas? Possibly still a bit clumsy if you can't capture I by putting everything in a class? Might still be a bit clumsy for add_xor_rot.

Is there a particular reason for the aversion to the current approach? It seems easier to read and probably also to debug to me.

I think this function should get a short description too. Wasn't immediately clear to me that it attempts to peel off larger blocks first before potentially finishing with a few smaller blocks.