authorgravatar for 124872+jedisct1@users.noreply.github.comFrank Denis <124872+jedisct1@users.noreply.github.com> 2023-05-22 20:33:35+02:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-05-22 20:33:35+02:00
log5af89b3dccf7ee375f68e9cd3ee4980fef89e38f
tree5262e6f92921170a2df1309f97d15669a7ffc614
parenteef92753c7cf677191adc40a7cdf7561ceb43bdb
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

std.crypto.chacha: support larger vectors on AVX2 and AVX512 targets (#15809)

* std.crypto.chacha: support larger vectors on AVX2 and AVX512 targets Ryzen 7 7700, ChaCha20/8 stream, long outputs: Generic: 3268 MiB/s AVX2 : 6023 MiB/s AVX512 : 8086 MiB/s Bump the rand.chacha buffer a tiny bit to take advantage of this. More than 8 blocks doesn't seem to make any measurable difference. ChaChaPoly also gets a small performance boost from this, albeit Poly1305 remains the bottleneck. Generic: 707 MiB/s AVX2 : 981 MiB/s AVX512 : 1202 MiB/s aarch64 appears to generally benefit from 4-way vectorization. Verified on Apple Silicon, but also on a Cortex A72.

2 files changed, 144 insertions(+), 52 deletions(-)

lib/std/crypto/chacha20.zig+143-51
......@@ -76,30 +76,98 @@ pub const XChaCha12Poly1305 = XChaChaPoly1305(12);
7676pub const XChaCha8Poly1305 = XChaChaPoly1305(8);
7777
7878// Vectorized implementation of the core function
79fn ChaChaVecImpl(comptime rounds_nb: usize) type {
79fn ChaChaVecImpl(comptime rounds_nb: usize, comptime degree: comptime_int) type {
8080 return struct {
81 const Lane = @Vector(4, u32);
81 const Lane = @Vector(4 * degree, u32);
8282 const BlockVec = [4]Lane;
8383
8484 fn initContext(key: [8]u32, d: [4]u32) BlockVec {
8585 const c = "expand 32-byte k";
86 const constant_le = comptime Lane{
87 mem.readIntLittle(u32, c[0..4]),
88 mem.readIntLittle(u32, c[4..8]),
89 mem.readIntLittle(u32, c[8..12]),
90 mem.readIntLittle(u32, c[12..16]),
91 };
92 return BlockVec{
93 constant_le,
94 Lane{ key[0], key[1], key[2], key[3] },
95 Lane{ key[4], key[5], key[6], key[7] },
96 Lane{ d[0], d[1], d[2], d[3] },
97 };
86 switch (degree) {
87 1 => {
88 const constant_le = Lane{
89 mem.readIntLittle(u32, c[0..4]),
90 mem.readIntLittle(u32, c[4..8]),
91 mem.readIntLittle(u32, c[8..12]),
92 mem.readIntLittle(u32, c[12..16]),
93 };
94 return BlockVec{
95 constant_le,
96 Lane{ key[0], key[1], key[2], key[3] },
97 Lane{ key[4], key[5], key[6], key[7] },
98 Lane{ d[0], d[1], d[2], d[3] },
99 };
100 },
101 2 => {
102 const constant_le = Lane{
103 mem.readIntLittle(u32, c[0..4]),
104 mem.readIntLittle(u32, c[4..8]),
105 mem.readIntLittle(u32, c[8..12]),
106 mem.readIntLittle(u32, c[12..16]),
107 mem.readIntLittle(u32, c[0..4]),
108 mem.readIntLittle(u32, c[4..8]),
109 mem.readIntLittle(u32, c[8..12]),
110 mem.readIntLittle(u32, c[12..16]),
111 };
112 return BlockVec{
113 constant_le,
114 Lane{ key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3] },
115 Lane{ key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7] },
116 Lane{ d[0], d[1], d[2], d[3], d[0] +% 1, d[1], d[2], d[3] },
117 };
118 },
119 4 => {
120 const constant_le = Lane{
121 mem.readIntLittle(u32, c[0..4]),
122 mem.readIntLittle(u32, c[4..8]),
123 mem.readIntLittle(u32, c[8..12]),
124 mem.readIntLittle(u32, c[12..16]),
125 mem.readIntLittle(u32, c[0..4]),
126 mem.readIntLittle(u32, c[4..8]),
127 mem.readIntLittle(u32, c[8..12]),
128 mem.readIntLittle(u32, c[12..16]),
129 mem.readIntLittle(u32, c[0..4]),
130 mem.readIntLittle(u32, c[4..8]),
131 mem.readIntLittle(u32, c[8..12]),
132 mem.readIntLittle(u32, c[12..16]),
133 mem.readIntLittle(u32, c[0..4]),
134 mem.readIntLittle(u32, c[4..8]),
135 mem.readIntLittle(u32, c[8..12]),
136 mem.readIntLittle(u32, c[12..16]),
137 };
138 return BlockVec{
139 constant_le,
140 Lane{ key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3] },
141 Lane{ key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7] },
142 Lane{ d[0], d[1], d[2], d[3], d[0] +% 1, d[1], d[2], d[3], d[0] +% 2, d[1], d[2], d[3], d[0] +% 3, d[1], d[2], d[3] },
143 };
144 },
145 else => @panic("invalid degree"),
146 }
98147 }
99148
100149 inline fn chacha20Core(x: *BlockVec, input: BlockVec) void {
101150 x.* = input;
102151
152 const m0 = switch (degree) {
153 1 => [_]i32{ 3, 0, 1, 2 },
154 2 => [_]i32{ 3, 0, 1, 2 } ++ [_]i32{ 7, 4, 5, 6 },
155 4 => [_]i32{ 3, 0, 1, 2 } ++ [_]i32{ 7, 4, 5, 6 } ++ [_]i32{ 11, 8, 9, 10 } ++ [_]i32{ 15, 12, 13, 14 },
156 else => @panic("invalid degree"),
157 };
158 const m1 = switch (degree) {
159 1 => [_]i32{ 2, 3, 0, 1 },
160 2 => [_]i32{ 2, 3, 0, 1 } ++ [_]i32{ 6, 7, 4, 5 },
161 4 => [_]i32{ 2, 3, 0, 1 } ++ [_]i32{ 6, 7, 4, 5 } ++ [_]i32{ 10, 11, 8, 9 } ++ [_]i32{ 14, 15, 12, 13 },
162 else => @panic("invalid degree"),
163 };
164 const m2 = switch (degree) {
165 1 => [_]i32{ 1, 2, 3, 0 },
166 2 => [_]i32{ 1, 2, 3, 0 } ++ [_]i32{ 5, 6, 7, 4 },
167 4 => [_]i32{ 1, 2, 3, 0 } ++ [_]i32{ 5, 6, 7, 4 } ++ [_]i32{ 9, 10, 11, 8 } ++ [_]i32{ 13, 14, 15, 12 },
168 else => @panic("invalid degree"),
169 };
170
103171 var r: usize = 0;
104172 while (r < rounds_nb) : (r += 2) {
105173 x[0] +%= x[1];
......@@ -112,13 +180,13 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
112180
113181 x[0] +%= x[1];
114182 x[3] ^= x[0];
115 x[0] = @shuffle(u32, x[0], undefined, [_]i32{ 3, 0, 1, 2 });
183 x[0] = @shuffle(u32, x[0], undefined, m0);
116184 x[3] = math.rotl(Lane, x[3], 8);
117185
118186 x[2] +%= x[3];
119 x[3] = @shuffle(u32, x[3], undefined, [_]i32{ 2, 3, 0, 1 });
187 x[3] = @shuffle(u32, x[3], undefined, m1);
120188 x[1] ^= x[2];
121 x[2] = @shuffle(u32, x[2], undefined, [_]i32{ 1, 2, 3, 0 });
189 x[2] = @shuffle(u32, x[2], undefined, m2);
122190 x[1] = math.rotl(Lane, x[1], 7);
123191
124192 x[0] +%= x[1];
......@@ -131,24 +199,26 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
131199
132200 x[0] +%= x[1];
133201 x[3] ^= x[0];
134 x[0] = @shuffle(u32, x[0], undefined, [_]i32{ 1, 2, 3, 0 });
202 x[0] = @shuffle(u32, x[0], undefined, m2);
135203 x[3] = math.rotl(Lane, x[3], 8);
136204
137205 x[2] +%= x[3];
138 x[3] = @shuffle(u32, x[3], undefined, [_]i32{ 2, 3, 0, 1 });
206 x[3] = @shuffle(u32, x[3], undefined, m1);
139207 x[1] ^= x[2];
140 x[2] = @shuffle(u32, x[2], undefined, [_]i32{ 3, 0, 1, 2 });
208 x[2] = @shuffle(u32, x[2], undefined, m0);
141209 x[1] = math.rotl(Lane, x[1], 7);
142210 }
143211 }
144212
145 inline fn hashToBytes(out: *[64]u8, x: BlockVec) void {
146 var i: usize = 0;
147 while (i < 4) : (i += 1) {
148 mem.writeIntLittle(u32, out[16 * i + 0 ..][0..4], x[i][0]);
149 mem.writeIntLittle(u32, out[16 * i + 4 ..][0..4], x[i][1]);
150 mem.writeIntLittle(u32, out[16 * i + 8 ..][0..4], x[i][2]);
151 mem.writeIntLittle(u32, out[16 * i + 12 ..][0..4], x[i][3]);
213 inline fn hashToBytes(comptime dm: usize, out: *[64 * dm]u8, x: BlockVec) void {
214 for (0..dm) |d| {
215 var i: usize = 0;
216 while (i < 4) : (i += 1) {
217 mem.writeIntLittle(u32, out[64 * d + 16 * i + 0 ..][0..4], x[i][0 + 4 * d]);
218 mem.writeIntLittle(u32, out[64 * d + 16 * i + 4 ..][0..4], x[i][1 + 4 * d]);
219 mem.writeIntLittle(u32, out[64 * d + 16 * i + 8 ..][0..4], x[i][2 + 4 * d]);
220 mem.writeIntLittle(u32, out[64 * d + 16 * i + 12 ..][0..4], x[i][3 + 4 * d]);
221 }
152222 }
153223 }
154224
......@@ -162,29 +232,33 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
162232 fn chacha20Xor(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
163233 var ctx = initContext(key, counter);
164234 var x: BlockVec = undefined;
165 var buf: [64]u8 = undefined;
235 var buf: [64 * degree]u8 = undefined;
166236 var i: usize = 0;
167 while (i + 64 <= in.len) : (i += 64) {
168 chacha20Core(x[0..], ctx);
169 contextFeedback(&x, ctx);
170 hashToBytes(buf[0..], x);
171
172 var xout = out[i..];
173 const xin = in[i..];
174 var j: usize = 0;
175 while (j < 64) : (j += 1) {
176 xout[j] = xin[j];
177 }
178 j = 0;
179 while (j < 64) : (j += 1) {
180 xout[j] ^= buf[j];
237 inline for ([_]comptime_int{ 4, 2, 1 }) |d| {
238 while (degree >= d and i + 64 * d <= in.len) : (i += 64 * d) {
239 chacha20Core(x[0..], ctx);
240 contextFeedback(&x, ctx);
241 hashToBytes(d, buf[0 .. 64 * d], x);
242
243 var xout = out[i..];
244 const xin = in[i..];
245 var j: usize = 0;
246 while (j < 64 * d) : (j += 1) {
247 xout[j] = xin[j];
248 }
249 j = 0;
250 while (j < 64 * d) : (j += 1) {
251 xout[j] ^= buf[j];
252 }
253 inline for (0..d) |d_| {
254 ctx[3][4 * d_] += @intCast(u32, d);
255 }
181256 }
182 ctx[3][0] += 1;
183257 }
184258 if (i < in.len) {
185259 chacha20Core(x[0..], ctx);
186260 contextFeedback(&x, ctx);
187 hashToBytes(buf[0..], x);
261 hashToBytes(1, buf[0..64], x);
188262
189263 var xout = out[i..];
190264 const xin = in[i..];
......@@ -199,18 +273,22 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
199273 var ctx = initContext(key, counter);
200274 var x: BlockVec = undefined;
201275 var i: usize = 0;
202 while (i + 64 <= out.len) : (i += 64) {
203 chacha20Core(x[0..], ctx);
204 contextFeedback(&x, ctx);
205 hashToBytes(out[i..][0..64], x);
206 ctx[3][0] += 1;
276 inline for ([_]comptime_int{ 4, 2, 1 }) |d| {
277 while (degree >= d and i + 64 * d <= out.len) : (i += 64 * d) {
278 chacha20Core(x[0..], ctx);
279 contextFeedback(&x, ctx);
280 hashToBytes(d, out[i..][0 .. 64 * d], x);
281 inline for (0..d) |d_| {
282 ctx[3][4 * d_] += @intCast(u32, d);
283 }
284 }
207285 }
208286 if (i < out.len) {
209287 chacha20Core(x[0..], ctx);
210288 contextFeedback(&x, ctx);
211289
212290 var buf: [64]u8 = undefined;
213 hashToBytes(buf[0..], x);
291 hashToBytes(1, buf[0..], x);
214292 @memcpy(out[i..], buf[0 .. out.len - i]);
215293 }
216294 }
......@@ -399,7 +477,21 @@ fn ChaChaNonVecImpl(comptime rounds_nb: usize) type {
399477}
400478
401479fn ChaChaImpl(comptime rounds_nb: usize) type {
402 return if (builtin.cpu.arch == .x86_64) ChaChaVecImpl(rounds_nb) else ChaChaNonVecImpl(rounds_nb);
480 switch (builtin.cpu.arch) {
481 .x86_64 => {
482 const has_avx2 = std.Target.x86.featureSetHas(builtin.cpu.features, .avx2);
483 const has_avx512f = std.Target.x86.featureSetHas(builtin.cpu.features, .avx512f);
484 if (has_avx512f) return ChaChaVecImpl(rounds_nb, 4);
485 if (has_avx2) return ChaChaVecImpl(rounds_nb, 2);
486 return ChaChaVecImpl(rounds_nb, 1);
487 },
488 .aarch64 => {
489 const has_neon = std.Target.aarch64.featureSetHas(builtin.cpu.features, .neon);
490 if (has_neon) return ChaChaVecImpl(rounds_nb, 4);
491 return ChaChaNonVecImpl(rounds_nb);
492 },
493 else => return ChaChaNonVecImpl(rounds_nb),
494 }
403495}
404496
405497fn keyToWords(key: [32]u8) [8]u32 {
lib/std/rand/ChaCha.zig+1-1
......@@ -10,7 +10,7 @@ const Self = @This();
1010
1111const Cipher = std.crypto.stream.chacha.ChaCha8IETF;
1212
13const State = [2 * Cipher.block_length]u8;
13const State = [8 * Cipher.block_length]u8;
1414
1515state: State,
1616offset: usize,