avx2.c (7793B)
1 #include <stddef.h> 2 #include <stdint.h> 3 4 #include "../keccak-fast-internal.h" 5 6 #if (defined(__x86_64__) || defined(_M_X64) || defined(__i386__)) && \ 7 !defined(KECCAK_FAST_NO_AVX2) 8 9 #include <immintrin.h> 10 11 static const uint64_t RC[24] = { 12 1ULL, 0x8082ULL, 0x800000000000808aULL, 0x8000000080008000ULL, 13 0x808bULL, 0x80000001ULL, 0x8000000080008081ULL, 0x8000000000008009ULL, 14 0x8aULL, 0x88ULL, 0x80008009ULL, 0x8000000aULL, 15 0x8000808bULL, 0x800000000000008bULL, 0x8000000000008089ULL, 0x8000000000008003ULL, 16 0x8000000000008002ULL, 0x8000000000000080ULL, 0x800aULL, 0x800000008000000aULL, 17 0x8000000080008081ULL, 0x8000000000008080ULL, 0x80000001ULL, 0x8000000080008008ULL}; 18 19 #define KF_ROL64(x, n) \ 20 _mm256_or_si256(_mm256_slli_epi64((x), (n)), _mm256_srli_epi64((x), 64 - (n))) 21 22 __attribute__((target("avx2"), always_inline)) 23 static inline void kf_avx2_round_x4(__m256i v[25], uint64_t rc) { 24 __m256i c0, c1, c2, c3, c4; 25 __m256i d0, d1, d2, d3, d4; 26 __m256i t, u; 27 28 #define KF_C(x) \ 29 _mm256_xor_si256(_mm256_xor_si256(v[x], v[(x) + 5]), \ 30 _mm256_xor_si256(_mm256_xor_si256(v[(x) + 10], v[(x) + 15]), \ 31 v[(x) + 20])) 32 c0 = KF_C(0); 33 c1 = KF_C(1); 34 c2 = KF_C(2); 35 c3 = KF_C(3); 36 c4 = KF_C(4); 37 #undef KF_C 38 39 d0 = _mm256_xor_si256(c4, KF_ROL64(c1, 1)); 40 d1 = _mm256_xor_si256(c0, KF_ROL64(c2, 1)); 41 d2 = _mm256_xor_si256(c1, KF_ROL64(c3, 1)); 42 d3 = _mm256_xor_si256(c2, KF_ROL64(c4, 1)); 43 d4 = _mm256_xor_si256(c3, KF_ROL64(c0, 1)); 44 45 #define KF_ADD(x, d) \ 46 v[(x)] = _mm256_xor_si256(v[(x)], (d)); \ 47 v[(x) + 5] = _mm256_xor_si256(v[(x) + 5], (d)); \ 48 v[(x) + 10] = _mm256_xor_si256(v[(x) + 10], (d)); \ 49 v[(x) + 15] = _mm256_xor_si256(v[(x) + 15], (d)); \ 50 v[(x) + 20] = _mm256_xor_si256(v[(x) + 20], (d)); 51 KF_ADD(0, d0) 52 KF_ADD(1, d1) 53 KF_ADD(2, d2) 54 KF_ADD(3, d3) 55 KF_ADD(4, d4) 56 #undef KF_ADD 57 58 t = v[1]; 59 #define KF_RHOPI4(dest, rot) \ 60 do { \ 61 u = v[dest]; \ 62 v[dest] = KF_ROL64(t, rot); \ 63 t = u; \ 64 } while (0) 65 KF_RHOPI4(10, 1); 66 KF_RHOPI4(7, 3); 67 KF_RHOPI4(11, 6); 68 KF_RHOPI4(17, 10); 69 KF_RHOPI4(18, 15); 70 KF_RHOPI4(3, 21); 71 KF_RHOPI4(5, 28); 72 KF_RHOPI4(16, 36); 73 KF_RHOPI4(8, 45); 74 KF_RHOPI4(21, 55); 75 KF_RHOPI4(24, 2); 76 KF_RHOPI4(4, 14); 77 KF_RHOPI4(15, 27); 78 KF_RHOPI4(23, 41); 79 KF_RHOPI4(19, 56); 80 KF_RHOPI4(13, 8); 81 KF_RHOPI4(12, 25); 82 KF_RHOPI4(2, 43); 83 KF_RHOPI4(20, 62); 84 KF_RHOPI4(14, 18); 85 KF_RHOPI4(22, 39); 86 KF_RHOPI4(9, 61); 87 KF_RHOPI4(6, 20); 88 KF_RHOPI4(1, 44); 89 #undef KF_RHOPI4 90 91 #define KF_ROW(y) \ 92 do { \ 93 __m256i r0 = v[(y)], r1 = v[(y) + 1], r2 = v[(y) + 2], r3 = v[(y) + 3], \ 94 r4 = v[(y) + 4]; \ 95 v[(y)] = _mm256_xor_si256(r0, _mm256_andnot_si256(r1, r2)); \ 96 v[(y) + 1] = _mm256_xor_si256(r1, _mm256_andnot_si256(r2, r3)); \ 97 v[(y) + 2] = _mm256_xor_si256(r2, _mm256_andnot_si256(r3, r4)); \ 98 v[(y) + 3] = _mm256_xor_si256(r3, _mm256_andnot_si256(r4, r0)); \ 99 v[(y) + 4] = _mm256_xor_si256(r4, _mm256_andnot_si256(r0, r1)); \ 100 } while (0) 101 KF_ROW(0); 102 KF_ROW(5); 103 KF_ROW(10); 104 KF_ROW(15); 105 KF_ROW(20); 106 #undef KF_ROW 107 108 v[0] = _mm256_xor_si256(v[0], _mm256_set1_epi64x((long long)rc)); 109 } 110 111 __attribute__((target("avx2"))) 112 static void kf_permute_x4(uint64_t s[4][25]) { 113 __m256i v[25]; 114 __m256i t[4]; 115 uint64_t tail[4]; 116 117 /* 4x4 transpose of 64-bit elements: 4 unpack + 4 lane permute. */ 118 #define KF_TRANSPOSE4(T) \ 119 do { \ 120 __m256i t0 = _mm256_unpacklo_epi64((T)[0], (T)[1]); \ 121 __m256i t1 = _mm256_unpackhi_epi64((T)[0], (T)[1]); \ 122 __m256i t2 = _mm256_unpacklo_epi64((T)[2], (T)[3]); \ 123 __m256i t3 = _mm256_unpackhi_epi64((T)[2], (T)[3]); \ 124 (T)[0] = _mm256_permute2x128_si256(t0, t2, 0x20); \ 125 (T)[1] = _mm256_permute2x128_si256(t1, t3, 0x20); \ 126 (T)[2] = _mm256_permute2x128_si256(t0, t2, 0x31); \ 127 (T)[3] = _mm256_permute2x128_si256(t1, t3, 0x31); \ 128 } while (0) 129 130 #define KF_LOAD4(B) \ 131 do { \ 132 for (int i = 0; i < 4; i++) { \ 133 t[i] = _mm256_loadu_si256((const void *)(s[i] + (B))); \ 134 } \ 135 KF_TRANSPOSE4(t); \ 136 for (int j = 0; j < 4; j++) v[(B) + j] = t[j]; \ 137 } while (0) 138 139 KF_LOAD4(0); 140 KF_LOAD4(4); 141 KF_LOAD4(8); 142 KF_LOAD4(12); 143 KF_LOAD4(16); 144 KF_LOAD4(20); 145 v[24] = _mm256_set_epi64x((long long)s[3][24], (long long)s[2][24], 146 (long long)s[1][24], (long long)s[0][24]); 147 148 #undef KF_LOAD4 149 150 for (int i = 0; i < 24; i++) kf_avx2_round_x4(v, RC[i]); 151 152 #define KF_STORE4(B) \ 153 do { \ 154 for (int j = 0; j < 4; j++) t[j] = v[(B) + j]; \ 155 KF_TRANSPOSE4(t); \ 156 for (int i = 0; i < 4; i++) { \ 157 _mm256_storeu_si256((void *)(s[i] + (B)), t[i]); \ 158 } \ 159 } while (0) 160 161 KF_STORE4(0); 162 KF_STORE4(4); 163 KF_STORE4(8); 164 KF_STORE4(12); 165 KF_STORE4(16); 166 KF_STORE4(20); 167 168 #undef KF_STORE4 169 #undef KF_TRANSPOSE4 170 171 _mm256_storeu_si256((void *)tail, v[24]); 172 for (int i = 0; i < 4; i++) s[i][24] = tail[i]; 173 } 174 175 __attribute__((target("avx2"))) 176 static void kf_permute_x4_12(uint64_t s[4][25]) { 177 __m256i v[25]; 178 uint64_t tmp[4]; 179 180 for (int j = 0; j < 25; j++) { 181 for (int i = 0; i < 4; i++) tmp[i] = s[i][j]; 182 v[j] = _mm256_loadu_si256((const __m256i *)tmp); 183 } 184 for (int i = 12; i < 24; i++) kf_avx2_round_x4(v, RC[i]); 185 for (int j = 0; j < 25; j++) { 186 _mm256_storeu_si256((__m256i *)tmp, v[j]); 187 for (int i = 0; i < 4; i++) s[i][j] = tmp[i]; 188 } 189 } 190 191 __attribute__((target("avx2"))) 192 void kf_avx2_batch(uint64_t states[][25], size_t count) { 193 size_t i = 0; 194 for (; i + 4 <= count; i += 4) { 195 kf_permute_x4(&states[i]); 196 } 197 for (; i < count; i++) { 198 kf_scalar_permute(states[i]); 199 } 200 } 201 202 __attribute__((target("avx2"))) 203 void kf_avx2_batch12(uint64_t states[][25], size_t count) { 204 size_t i = 0; 205 for (; i + 4 <= count; i += 4) { 206 kf_permute_x4_12(&states[i]); 207 } 208 if (i < count) kf_scalar_batch12(&states[i], count - i); 209 } 210 211 __attribute__((constructor)) static void kf_avx2_register(void) { 212 if (!__builtin_cpu_supports("avx2")) { 213 return; 214 } 215 if (kf_batch_priority <= KF_PRIO_AVX2) { 216 kf_batch_perm = kf_avx2_batch; 217 kf_batch_perm12 = kf_avx2_batch12; 218 kf_batch_priority = KF_PRIO_AVX2; 219 kf_batch_lanes = 4; 220 kf_batch_label = "avx2"; 221 } 222 } 223 224 #endif