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);...@@ -76,30 +76,98 @@ pub const XChaCha12Poly1305 = XChaChaPoly1305(12);
76pub const XChaCha8Poly1305 = XChaChaPoly1305(8);76pub const XChaCha8Poly1305 = XChaChaPoly1305(8);
7777
78// Vectorized implementation of the core function78// Vectorized implementation of the core function
79fn ChaChaVecImpl(comptime rounds_nb: usize) type {79fn ChaChaVecImpl(comptime rounds_nb: usize, comptime degree: comptime_int) type {
80 return struct {80 return struct {
81 const Lane = @Vector(4, u32);81 const Lane = @Vector(4 * degree, u32);
82 const BlockVec = [4]Lane;82 const BlockVec = [4]Lane;
8383
84 fn initContext(key: [8]u32, d: [4]u32) BlockVec {84 fn initContext(key: [8]u32, d: [4]u32) BlockVec {
85 const c = "expand 32-byte k";85 const c = "expand 32-byte k";
86 const constant_le = comptime Lane{86 switch (degree) {
87 mem.readIntLittle(u32, c[0..4]),87 1 => {
88 mem.readIntLittle(u32, c[4..8]),88 const constant_le = Lane{
89 mem.readIntLittle(u32, c[8..12]),89 mem.readIntLittle(u32, c[0..4]),
90 mem.readIntLittle(u32, c[12..16]),90 mem.readIntLittle(u32, c[4..8]),
91 };91 mem.readIntLittle(u32, c[8..12]),
92 return BlockVec{92 mem.readIntLittle(u32, c[12..16]),
93 constant_le,93 };
94 Lane{ key[0], key[1], key[2], key[3] },94 return BlockVec{
95 Lane{ key[4], key[5], key[6], key[7] },95 constant_le,
96 Lane{ d[0], d[1], d[2], d[3] },96 Lane{ key[0], key[1], key[2], key[3] },
97 };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 }
98 }147 }
99148
100 inline fn chacha20Core(x: *BlockVec, input: BlockVec) void {149 inline fn chacha20Core(x: *BlockVec, input: BlockVec) void {
101 x.* = input;150 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
103 var r: usize = 0;171 var r: usize = 0;
104 while (r < rounds_nb) : (r += 2) {172 while (r < rounds_nb) : (r += 2) {
105 x[0] +%= x[1];173 x[0] +%= x[1];
...@@ -112,13 +180,13 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {...@@ -112,13 +180,13 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
112180
113 x[0] +%= x[1];181 x[0] +%= x[1];
114 x[3] ^= x[0];182 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);
116 x[3] = math.rotl(Lane, x[3], 8);184 x[3] = math.rotl(Lane, x[3], 8);
117185
118 x[2] +%= x[3];186 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);
120 x[1] ^= x[2];188 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);
122 x[1] = math.rotl(Lane, x[1], 7);190 x[1] = math.rotl(Lane, x[1], 7);
123191
124 x[0] +%= x[1];192 x[0] +%= x[1];
...@@ -131,24 +199,26 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {...@@ -131,24 +199,26 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
131199
132 x[0] +%= x[1];200 x[0] +%= x[1];
133 x[3] ^= x[0];201 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);
135 x[3] = math.rotl(Lane, x[3], 8);203 x[3] = math.rotl(Lane, x[3], 8);
136204
137 x[2] +%= x[3];205 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);
139 x[1] ^= x[2];207 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);
141 x[1] = math.rotl(Lane, x[1], 7);209 x[1] = math.rotl(Lane, x[1], 7);
142 }210 }
143 }211 }
144212
145 inline fn hashToBytes(out: *[64]u8, x: BlockVec) void {213 inline fn hashToBytes(comptime dm: usize, out: *[64 * dm]u8, x: BlockVec) void {
146 var i: usize = 0;214 for (0..dm) |d| {
147 while (i < 4) : (i += 1) {215 var i: usize = 0;
148 mem.writeIntLittle(u32, out[16 * i + 0 ..][0..4], x[i][0]);216 while (i < 4) : (i += 1) {
149 mem.writeIntLittle(u32, out[16 * i + 4 ..][0..4], x[i][1]);217 mem.writeIntLittle(u32, out[64 * d + 16 * i + 0 ..][0..4], x[i][0 + 4 * d]);
150 mem.writeIntLittle(u32, out[16 * i + 8 ..][0..4], x[i][2]);218 mem.writeIntLittle(u32, out[64 * d + 16 * i + 4 ..][0..4], x[i][1 + 4 * d]);
151 mem.writeIntLittle(u32, out[16 * i + 12 ..][0..4], x[i][3]);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 }
152 }222 }
153 }223 }
154224
...@@ -162,29 +232,33 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {...@@ -162,29 +232,33 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
162 fn chacha20Xor(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {232 fn chacha20Xor(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
163 var ctx = initContext(key, counter);233 var ctx = initContext(key, counter);
164 var x: BlockVec = undefined;234 var x: BlockVec = undefined;
165 var buf: [64]u8 = undefined;235 var buf: [64 * degree]u8 = undefined;
166 var i: usize = 0;236 var i: usize = 0;
167 while (i + 64 <= in.len) : (i += 64) {237 inline for ([_]comptime_int{ 4, 2, 1 }) |d| {
168 chacha20Core(x[0..], ctx);238 while (degree >= d and i + 64 * d <= in.len) : (i += 64 * d) {
169 contextFeedback(&x, ctx);239 chacha20Core(x[0..], ctx);
170 hashToBytes(buf[0..], x);240 contextFeedback(&x, ctx);
171241 hashToBytes(d, buf[0 .. 64 * d], x);
172 var xout = out[i..];242
173 const xin = in[i..];243 var xout = out[i..];
174 var j: usize = 0;244 const xin = in[i..];
175 while (j < 64) : (j += 1) {245 var j: usize = 0;
176 xout[j] = xin[j];246 while (j < 64 * d) : (j += 1) {
177 }247 xout[j] = xin[j];
178 j = 0;248 }
179 while (j < 64) : (j += 1) {249 j = 0;
180 xout[j] ^= buf[j];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 }
181 }256 }
182 ctx[3][0] += 1;
183 }257 }
184 if (i < in.len) {258 if (i < in.len) {
185 chacha20Core(x[0..], ctx);259 chacha20Core(x[0..], ctx);
186 contextFeedback(&x, ctx);260 contextFeedback(&x, ctx);
187 hashToBytes(buf[0..], x);261 hashToBytes(1, buf[0..64], x);
188262
189 var xout = out[i..];263 var xout = out[i..];
190 const xin = in[i..];264 const xin = in[i..];
...@@ -199,18 +273,22 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {...@@ -199,18 +273,22 @@ fn ChaChaVecImpl(comptime rounds_nb: usize) type {
199 var ctx = initContext(key, counter);273 var ctx = initContext(key, counter);
200 var x: BlockVec = undefined;274 var x: BlockVec = undefined;
201 var i: usize = 0;275 var i: usize = 0;
202 while (i + 64 <= out.len) : (i += 64) {276 inline for ([_]comptime_int{ 4, 2, 1 }) |d| {
203 chacha20Core(x[0..], ctx);277 while (degree >= d and i + 64 * d <= out.len) : (i += 64 * d) {
204 contextFeedback(&x, ctx);278 chacha20Core(x[0..], ctx);
205 hashToBytes(out[i..][0..64], x);279 contextFeedback(&x, ctx);
206 ctx[3][0] += 1;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 }
207 }285 }
208 if (i < out.len) {286 if (i < out.len) {
209 chacha20Core(x[0..], ctx);287 chacha20Core(x[0..], ctx);
210 contextFeedback(&x, ctx);288 contextFeedback(&x, ctx);
211289
212 var buf: [64]u8 = undefined;290 var buf: [64]u8 = undefined;
213 hashToBytes(buf[0..], x);291 hashToBytes(1, buf[0..], x);
214 @memcpy(out[i..], buf[0 .. out.len - i]);292 @memcpy(out[i..], buf[0 .. out.len - i]);
215 }293 }
216 }294 }
...@@ -399,7 +477,21 @@ fn ChaChaNonVecImpl(comptime rounds_nb: usize) type {...@@ -399,7 +477,21 @@ fn ChaChaNonVecImpl(comptime rounds_nb: usize) type {
399}477}
400478
401fn ChaChaImpl(comptime rounds_nb: usize) type {479fn 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 }
403}495}
404496
405fn keyToWords(key: [32]u8) [8]u32 {497fn keyToWords(key: [32]u8) [8]u32 {
lib/std/rand/ChaCha.zig+1-1
...@@ -10,7 +10,7 @@ const Self = @This();...@@ -10,7 +10,7 @@ const Self = @This();
1010
11const Cipher = std.crypto.stream.chacha.ChaCha8IETF;11const Cipher = std.crypto.stream.chacha.ChaCha8IETF;
1212
13const State = [2 * Cipher.block_length]u8;13const State = [8 * Cipher.block_length]u8;
1414
15state: State,15state: State,
16offset: usize,16offset: usize,