commit 32859a24f992a46fef4009946b7a7fd287034938
parent 1d186f9effe31f695b67b5627edd4dc6e9d0da59
Author: finwo <finwo@pm.me>
Date: Sat, 10 Oct 2026 22:50:54 +0200
transpose 4x25/8x25 state in vector blocks instead of scalar gather
Diffstat:
3 files changed, 132 insertions(+), 18 deletions(-)
diff --git a/README.md b/README.md
@@ -51,6 +51,12 @@ locals, fuses theta into the previous round's chi, and unrolls the round loop
permutation with `bmi` enabled, which turns chi's `not`+`and` into `andn`; gcc
also gets `rorx` for the rotates. About 20% faster than the donor.
+The AVX2 and AVX-512 batch kernels hold the 4 or 8 messages of a group one per
+vector lane, so the 25 state words live in 25 vector registers. The state is
+transposed in 8x8 (or 4x4) blocks with unpack and shuffle instructions at the
+sponge boundary. Their round bodies unroll theta and chi by hand so `-O2` emits
+the same code as `-O3`.
+
## Dev
```
diff --git a/src/backend/avx2.c b/src/backend/avx2.c
@@ -110,17 +110,65 @@ static inline void kf_avx2_round_x4(__m256i v[25], uint64_t rc) {
__attribute__((target("avx2")))
static void kf_permute_x4(uint64_t s[4][25]) {
__m256i v[25];
- uint64_t tmp[4];
+ __m256i t[4];
+ uint64_t tail[4];
+
+/* 4x4 transpose of 64-bit elements: 4 unpack + 4 lane permute. */
+#define KF_TRANSPOSE4(T) \
+ do { \
+ __m256i t0 = _mm256_unpacklo_epi64((T)[0], (T)[1]); \
+ __m256i t1 = _mm256_unpackhi_epi64((T)[0], (T)[1]); \
+ __m256i t2 = _mm256_unpacklo_epi64((T)[2], (T)[3]); \
+ __m256i t3 = _mm256_unpackhi_epi64((T)[2], (T)[3]); \
+ (T)[0] = _mm256_permute2x128_si256(t0, t2, 0x20); \
+ (T)[1] = _mm256_permute2x128_si256(t1, t3, 0x20); \
+ (T)[2] = _mm256_permute2x128_si256(t0, t2, 0x31); \
+ (T)[3] = _mm256_permute2x128_si256(t1, t3, 0x31); \
+ } while (0)
+
+#define KF_LOAD4(B) \
+ do { \
+ for (int i = 0; i < 4; i++) { \
+ t[i] = _mm256_loadu_si256((const void *)(s[i] + (B))); \
+ } \
+ KF_TRANSPOSE4(t); \
+ for (int j = 0; j < 4; j++) v[(B) + j] = t[j]; \
+ } while (0)
+
+ KF_LOAD4(0);
+ KF_LOAD4(4);
+ KF_LOAD4(8);
+ KF_LOAD4(12);
+ KF_LOAD4(16);
+ KF_LOAD4(20);
+ v[24] = _mm256_set_epi64x((long long)s[3][24], (long long)s[2][24],
+ (long long)s[1][24], (long long)s[0][24]);
+
+#undef KF_LOAD4
- for (int j = 0; j < 25; j++) {
- for (int i = 0; i < 4; i++) tmp[i] = s[i][j];
- v[j] = _mm256_loadu_si256((const __m256i *)tmp);
- }
for (int i = 0; i < 24; i++) kf_avx2_round_x4(v, RC[i]);
- for (int j = 0; j < 25; j++) {
- _mm256_storeu_si256((__m256i *)tmp, v[j]);
- for (int i = 0; i < 4; i++) s[i][j] = tmp[i];
- }
+
+#define KF_STORE4(B) \
+ do { \
+ for (int j = 0; j < 4; j++) t[j] = v[(B) + j]; \
+ KF_TRANSPOSE4(t); \
+ for (int i = 0; i < 4; i++) { \
+ _mm256_storeu_si256((void *)(s[i] + (B)), t[i]); \
+ } \
+ } while (0)
+
+ KF_STORE4(0);
+ KF_STORE4(4);
+ KF_STORE4(8);
+ KF_STORE4(12);
+ KF_STORE4(16);
+ KF_STORE4(20);
+
+#undef KF_STORE4
+#undef KF_TRANSPOSE4
+
+ _mm256_storeu_si256((void *)tail, v[24]);
+ for (int i = 0; i < 4; i++) s[i][24] = tail[i];
}
__attribute__((target("avx2")))
diff --git a/src/backend/avx512.c b/src/backend/avx512.c
@@ -107,17 +107,77 @@ static inline void kf_avx512_round_x8(__m512i v[25], uint64_t rc) {
__attribute__((target("avx512f")))
static void kf_permute_x8(uint64_t s[8][25]) {
__m512i v[25];
- uint64_t tmp[8];
+ __m512i t[8];
+ uint64_t tail[8];
+
+/* 8x8 transpose of 64-bit elements: 8 unpack + 16 shuffle. */
+#define KF_TRANSPOSE8(T) \
+ do { \
+ __m512i t0 = _mm512_unpacklo_epi64((T)[0], (T)[1]); \
+ __m512i t1 = _mm512_unpackhi_epi64((T)[0], (T)[1]); \
+ __m512i t2 = _mm512_unpacklo_epi64((T)[2], (T)[3]); \
+ __m512i t3 = _mm512_unpackhi_epi64((T)[2], (T)[3]); \
+ __m512i t4 = _mm512_unpacklo_epi64((T)[4], (T)[5]); \
+ __m512i t5 = _mm512_unpackhi_epi64((T)[4], (T)[5]); \
+ __m512i t6 = _mm512_unpacklo_epi64((T)[6], (T)[7]); \
+ __m512i t7 = _mm512_unpackhi_epi64((T)[6], (T)[7]); \
+ __m512i s0 = _mm512_shuffle_i64x2(t0, t2, 0x88); \
+ __m512i s1 = _mm512_shuffle_i64x2(t1, t3, 0x88); \
+ __m512i s2 = _mm512_shuffle_i64x2(t0, t2, 0xdd); \
+ __m512i s3 = _mm512_shuffle_i64x2(t1, t3, 0xdd); \
+ __m512i s4 = _mm512_shuffle_i64x2(t4, t6, 0x88); \
+ __m512i s5 = _mm512_shuffle_i64x2(t5, t7, 0x88); \
+ __m512i s6 = _mm512_shuffle_i64x2(t4, t6, 0xdd); \
+ __m512i s7 = _mm512_shuffle_i64x2(t5, t7, 0xdd); \
+ (T)[0] = _mm512_shuffle_i64x2(s0, s4, 0x88); \
+ (T)[1] = _mm512_shuffle_i64x2(s1, s5, 0x88); \
+ (T)[2] = _mm512_shuffle_i64x2(s2, s6, 0x88); \
+ (T)[3] = _mm512_shuffle_i64x2(s3, s7, 0x88); \
+ (T)[4] = _mm512_shuffle_i64x2(s0, s4, 0xdd); \
+ (T)[5] = _mm512_shuffle_i64x2(s1, s5, 0xdd); \
+ (T)[6] = _mm512_shuffle_i64x2(s2, s6, 0xdd); \
+ (T)[7] = _mm512_shuffle_i64x2(s3, s7, 0xdd); \
+ } while (0)
+
+#define KF_LOAD8(B) \
+ do { \
+ for (int i = 0; i < 8; i++) { \
+ t[i] = _mm512_loadu_si512((const void *)(s[i] + (B))); \
+ } \
+ KF_TRANSPOSE8(t); \
+ for (int j = 0; j < 8; j++) v[(B) + j] = t[j]; \
+ } while (0)
+
+ KF_LOAD8(0);
+ KF_LOAD8(8);
+ KF_LOAD8(16);
+ v[24] = _mm512_set_epi64((long long)s[7][24], (long long)s[6][24],
+ (long long)s[5][24], (long long)s[4][24],
+ (long long)s[3][24], (long long)s[2][24],
+ (long long)s[1][24], (long long)s[0][24]);
+
+#undef KF_LOAD8
- for (int j = 0; j < 25; j++) {
- for (int i = 0; i < 8; i++) tmp[i] = s[i][j];
- v[j] = _mm512_loadu_si512((const __m512i *)tmp);
- }
for (int i = 0; i < 24; i++) kf_avx512_round_x8(v, RC[i]);
- for (int j = 0; j < 25; j++) {
- _mm512_storeu_si512((__m512i *)tmp, v[j]);
- for (int i = 0; i < 8; i++) s[i][j] = tmp[i];
- }
+
+#define KF_STORE8(B) \
+ do { \
+ for (int j = 0; j < 8; j++) t[j] = v[(B) + j]; \
+ KF_TRANSPOSE8(t); \
+ for (int i = 0; i < 8; i++) { \
+ _mm512_storeu_si512((void *)(s[i] + (B)), t[i]); \
+ } \
+ } while (0)
+
+ KF_STORE8(0);
+ KF_STORE8(8);
+ KF_STORE8(16);
+
+#undef KF_STORE8
+#undef KF_TRANSPOSE8
+
+ _mm512_storeu_si512((void *)tail, v[24]);
+ for (int i = 0; i < 8; i++) s[i][24] = tail[i];
}
__attribute__((target("avx512f")))
