avx512.c (9138B)
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_AVX512) 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 __attribute__((target("avx512f"), always_inline)) 20 static inline void kf_avx512_round_x8(__m512i v[25], uint64_t rc) { 21 __m512i c0, c1, c2, c3, c4; 22 __m512i d0, d1, d2, d3, d4; 23 __m512i t, u; 24 25 #define KF_C(x) \ 26 _mm512_ternarylogic_epi64( \ 27 _mm512_ternarylogic_epi64(v[x], v[(x) + 5], v[(x) + 10], 0x96), \ 28 v[(x) + 15], v[(x) + 20], 0x96) 29 c0 = KF_C(0); 30 c1 = KF_C(1); 31 c2 = KF_C(2); 32 c3 = KF_C(3); 33 c4 = KF_C(4); 34 #undef KF_C 35 36 d0 = _mm512_xor_si512(c4, _mm512_rol_epi64(c1, 1)); 37 d1 = _mm512_xor_si512(c0, _mm512_rol_epi64(c2, 1)); 38 d2 = _mm512_xor_si512(c1, _mm512_rol_epi64(c3, 1)); 39 d3 = _mm512_xor_si512(c2, _mm512_rol_epi64(c4, 1)); 40 d4 = _mm512_xor_si512(c3, _mm512_rol_epi64(c0, 1)); 41 42 #define KF_ADD(x, d) \ 43 v[(x)] = _mm512_xor_si512(v[(x)], (d)); \ 44 v[(x) + 5] = _mm512_xor_si512(v[(x) + 5], (d)); \ 45 v[(x) + 10] = _mm512_xor_si512(v[(x) + 10], (d)); \ 46 v[(x) + 15] = _mm512_xor_si512(v[(x) + 15], (d)); \ 47 v[(x) + 20] = _mm512_xor_si512(v[(x) + 20], (d)); 48 KF_ADD(0, d0) 49 KF_ADD(1, d1) 50 KF_ADD(2, d2) 51 KF_ADD(3, d3) 52 KF_ADD(4, d4) 53 #undef KF_ADD 54 55 t = v[1]; 56 #define KF_RHOPI8(dest, rot) \ 57 do { \ 58 u = v[dest]; \ 59 v[dest] = _mm512_rol_epi64(t, rot); \ 60 t = u; \ 61 } while (0) 62 KF_RHOPI8(10, 1); 63 KF_RHOPI8(7, 3); 64 KF_RHOPI8(11, 6); 65 KF_RHOPI8(17, 10); 66 KF_RHOPI8(18, 15); 67 KF_RHOPI8(3, 21); 68 KF_RHOPI8(5, 28); 69 KF_RHOPI8(16, 36); 70 KF_RHOPI8(8, 45); 71 KF_RHOPI8(21, 55); 72 KF_RHOPI8(24, 2); 73 KF_RHOPI8(4, 14); 74 KF_RHOPI8(15, 27); 75 KF_RHOPI8(23, 41); 76 KF_RHOPI8(19, 56); 77 KF_RHOPI8(13, 8); 78 KF_RHOPI8(12, 25); 79 KF_RHOPI8(2, 43); 80 KF_RHOPI8(20, 62); 81 KF_RHOPI8(14, 18); 82 KF_RHOPI8(22, 39); 83 KF_RHOPI8(9, 61); 84 KF_RHOPI8(6, 20); 85 KF_RHOPI8(1, 44); 86 #undef KF_RHOPI8 87 88 #define KF_ROW(y) \ 89 do { \ 90 __m512i r0 = v[(y)], r1 = v[(y) + 1], r2 = v[(y) + 2], r3 = v[(y) + 3], \ 91 r4 = v[(y) + 4]; \ 92 v[(y)] = _mm512_ternarylogic_epi64(r0, r1, r2, 0xD2); \ 93 v[(y) + 1] = _mm512_ternarylogic_epi64(r1, r2, r3, 0xD2); \ 94 v[(y) + 2] = _mm512_ternarylogic_epi64(r2, r3, r4, 0xD2); \ 95 v[(y) + 3] = _mm512_ternarylogic_epi64(r3, r4, r0, 0xD2); \ 96 v[(y) + 4] = _mm512_ternarylogic_epi64(r4, r0, r1, 0xD2); \ 97 } while (0) 98 KF_ROW(0); 99 KF_ROW(5); 100 KF_ROW(10); 101 KF_ROW(15); 102 KF_ROW(20); 103 #undef KF_ROW 104 105 v[0] = _mm512_xor_si512(v[0], _mm512_set1_epi64((long long)rc)); 106 } 107 108 __attribute__((target("avx512f"))) 109 static void kf_permute_x8(uint64_t s[8][25]) { 110 __m512i v[25]; 111 __m512i t[8]; 112 uint64_t tail[8]; 113 114 /* 8x8 transpose of 64-bit elements: 8 unpack + 16 shuffle. */ 115 #define KF_TRANSPOSE8(T) \ 116 do { \ 117 __m512i t0 = _mm512_unpacklo_epi64((T)[0], (T)[1]); \ 118 __m512i t1 = _mm512_unpackhi_epi64((T)[0], (T)[1]); \ 119 __m512i t2 = _mm512_unpacklo_epi64((T)[2], (T)[3]); \ 120 __m512i t3 = _mm512_unpackhi_epi64((T)[2], (T)[3]); \ 121 __m512i t4 = _mm512_unpacklo_epi64((T)[4], (T)[5]); \ 122 __m512i t5 = _mm512_unpackhi_epi64((T)[4], (T)[5]); \ 123 __m512i t6 = _mm512_unpacklo_epi64((T)[6], (T)[7]); \ 124 __m512i t7 = _mm512_unpackhi_epi64((T)[6], (T)[7]); \ 125 __m512i s0 = _mm512_shuffle_i64x2(t0, t2, 0x88); \ 126 __m512i s1 = _mm512_shuffle_i64x2(t1, t3, 0x88); \ 127 __m512i s2 = _mm512_shuffle_i64x2(t0, t2, 0xdd); \ 128 __m512i s3 = _mm512_shuffle_i64x2(t1, t3, 0xdd); \ 129 __m512i s4 = _mm512_shuffle_i64x2(t4, t6, 0x88); \ 130 __m512i s5 = _mm512_shuffle_i64x2(t5, t7, 0x88); \ 131 __m512i s6 = _mm512_shuffle_i64x2(t4, t6, 0xdd); \ 132 __m512i s7 = _mm512_shuffle_i64x2(t5, t7, 0xdd); \ 133 (T)[0] = _mm512_shuffle_i64x2(s0, s4, 0x88); \ 134 (T)[1] = _mm512_shuffle_i64x2(s1, s5, 0x88); \ 135 (T)[2] = _mm512_shuffle_i64x2(s2, s6, 0x88); \ 136 (T)[3] = _mm512_shuffle_i64x2(s3, s7, 0x88); \ 137 (T)[4] = _mm512_shuffle_i64x2(s0, s4, 0xdd); \ 138 (T)[5] = _mm512_shuffle_i64x2(s1, s5, 0xdd); \ 139 (T)[6] = _mm512_shuffle_i64x2(s2, s6, 0xdd); \ 140 (T)[7] = _mm512_shuffle_i64x2(s3, s7, 0xdd); \ 141 } while (0) 142 143 #define KF_LOAD8(B) \ 144 do { \ 145 for (int i = 0; i < 8; i++) { \ 146 t[i] = _mm512_loadu_si512((const void *)(s[i] + (B))); \ 147 } \ 148 KF_TRANSPOSE8(t); \ 149 for (int j = 0; j < 8; j++) v[(B) + j] = t[j]; \ 150 } while (0) 151 152 KF_LOAD8(0); 153 KF_LOAD8(8); 154 KF_LOAD8(16); 155 v[24] = _mm512_set_epi64((long long)s[7][24], (long long)s[6][24], 156 (long long)s[5][24], (long long)s[4][24], 157 (long long)s[3][24], (long long)s[2][24], 158 (long long)s[1][24], (long long)s[0][24]); 159 160 #undef KF_LOAD8 161 162 for (int i = 0; i < 24; i++) kf_avx512_round_x8(v, RC[i]); 163 164 #define KF_STORE8(B) \ 165 do { \ 166 for (int j = 0; j < 8; j++) t[j] = v[(B) + j]; \ 167 KF_TRANSPOSE8(t); \ 168 for (int i = 0; i < 8; i++) { \ 169 _mm512_storeu_si512((void *)(s[i] + (B)), t[i]); \ 170 } \ 171 } while (0) 172 173 KF_STORE8(0); 174 KF_STORE8(8); 175 KF_STORE8(16); 176 177 #undef KF_STORE8 178 #undef KF_TRANSPOSE8 179 180 _mm512_storeu_si512((void *)tail, v[24]); 181 for (int i = 0; i < 8; i++) s[i][24] = tail[i]; 182 } 183 184 __attribute__((target("avx512f"))) 185 static void kf_permute_x8_12(uint64_t s[8][25]) { 186 __m512i v[25]; 187 uint64_t tmp[8]; 188 189 for (int j = 0; j < 25; j++) { 190 for (int i = 0; i < 8; i++) tmp[i] = s[i][j]; 191 v[j] = _mm512_loadu_si512((const __m512i *)tmp); 192 } 193 for (int i = 12; i < 24; i++) kf_avx512_round_x8(v, RC[i]); 194 for (int j = 0; j < 25; j++) { 195 _mm512_storeu_si512((__m512i *)tmp, v[j]); 196 for (int i = 0; i < 8; i++) s[i][j] = tmp[i]; 197 } 198 } 199 200 __attribute__((target("avx512f"))) 201 void kf_avx512_batch(uint64_t states[][25], size_t count) { 202 size_t i = 0; 203 for (; i + 8 <= count; i += 8) { 204 kf_permute_x8(&states[i]); 205 } 206 for (; i < count; i++) { 207 kf_scalar_permute(states[i]); 208 } 209 } 210 211 __attribute__((target("avx512f"))) 212 void kf_avx512_batch12(uint64_t states[][25], size_t count) { 213 size_t i = 0; 214 for (; i + 8 <= count; i += 8) { 215 kf_permute_x8_12(&states[i]); 216 } 217 if (i < count) kf_scalar_batch12(&states[i], count - i); 218 } 219 220 __attribute__((constructor)) static void kf_avx512_register(void) { 221 if (!__builtin_cpu_supports("avx512f")) { 222 return; 223 } 224 if (kf_batch_priority <= KF_PRIO_AVX512) { 225 kf_batch_perm = kf_avx512_batch; 226 kf_batch_perm12 = kf_avx512_batch12; 227 kf_batch_priority = KF_PRIO_AVX512; 228 kf_batch_lanes = 8; 229 kf_batch_label = "avx512"; 230 } 231 } 232 233 #endif