authorgravatar for 124872+jedisct1@users.noreply.github.comFrank Denis <124872+jedisct1@users.noreply.github.com> 2020-10-09 23:19:27+02:00
committergravatar for 124872+jedisct1@users.noreply.github.comFrank Denis <124872+jedisct1@users.noreply.github.com> 2020-10-10 22:45:41+02:00
log9b386bda33e94c79d6b9a1db911d394c26592e71
treef9d2a50e5664776a8c18dd71d6cd18bccd3cd135
parent53c63bdb73d9fbc5a54afb4977bb975b03c4c9cc

std/crypto: add a vectorized ChaCha20 implementation

Brings a 30% speed boost on x86_64 even though we still process only one block at a time for now. Only enabled on x86_64 since the non-vectorized implementation seems to currently perform better on some architectures (at least on aarch64). But the non-vectorized implementation still gets a little speed boost as well (~17%) with these changes.

1 files changed, 292 insertions(+), 118 deletions(-)

lib/std/crypto/chacha20.zig+292-118
...@@ -10,120 +10,315 @@ const mem = std.mem;...@@ -10,120 +10,315 @@ const mem = std.mem;
10const assert = std.debug.assert;10const assert = std.debug.assert;
11const testing = std.testing;11const testing = std.testing;
12const maxInt = std.math.maxInt;12const maxInt = std.math.maxInt;
13const Vector = std.meta.Vector;
13const Poly1305 = std.crypto.onetimeauth.Poly1305;14const Poly1305 = std.crypto.onetimeauth.Poly1305;
1415
15const QuarterRound = struct {16// Vectorized implementation of the core function
16 a: usize,17const ChaCha20VecImpl = struct {
17 b: usize,18 const Lane = Vector(4, u32);
18 c: usize,19 const BlockVec = [4]Lane;
19 d: usize,20
20};21 fn initContext(key: [8]u32, d: [4]u32) BlockVec {
22 const c = "expand 32-byte k";
23 const constant_le = comptime Lane{
24 mem.readIntLittle(u32, c[0..4]),
25 mem.readIntLittle(u32, c[4..8]),
26 mem.readIntLittle(u32, c[8..12]),
27 mem.readIntLittle(u32, c[12..16]),
28 };
29 return BlockVec{
30 constant_le,
31 Lane{ key[0], key[1], key[2], key[3] },
32 Lane{ key[4], key[5], key[6], key[7] },
33 Lane{ d[0], d[1], d[2], d[3] },
34 };
35 }
2136
22fn Rp(a: usize, b: usize, c: usize, d: usize) QuarterRound {37 inline fn chacha20Core(x: *BlockVec, input: BlockVec) void {
23 return QuarterRound{38 const rot8 = Vector(16, i32){ 3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14 };
24 .a = a,39 const rot16 = Vector(16, i32){ 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13 };
25 .b = b,40
26 .c = c,41 x.* = input;
27 .d = d,42
28 };43 var r: usize = 0;
29}44 while (r < 20) : (r += 2) {
45 x[0] +%= x[1];
46 x[3] ^= x[0];
47 x[3] = @bitCast(Vector(4, u32), @shuffle(u8, @bitCast(Vector(16, u8), x[3]), undefined, rot16));
48
49 x[2] +%= x[3];
50 x[1] ^= x[2];
51
52 var t1 = x[1];
53 x[1] <<= @splat(4, @as(u5, 12));
54 t1 >>= @splat(4, @as(u5, 20));
55 x[1] ^= t1;
56
57 x[0] +%= x[1];
58 x[3] ^= x[0];
59 x[0] = @shuffle(u32, x[0], undefined, Vector(4, i32){ 3, 0, 1, 2 });
60 x[3] = @bitCast(Vector(4, u32), @shuffle(u8, @bitCast(Vector(16, u8), x[3]), undefined, rot8));
61
62 x[2] +%= x[3];
63 x[3] = @shuffle(u32, x[3], undefined, Vector(4, i32){ 2, 3, 0, 1 });
64 x[1] ^= x[2];
65 x[2] = @shuffle(u32, x[2], undefined, Vector(4, i32){ 1, 2, 3, 0 });
66
67 t1 = x[1];
68 x[1] <<= @splat(4, @as(u5, 7));
69 t1 >>= @splat(4, @as(u5, 25));
70 x[1] ^= t1;
71
72 x[0] +%= x[1];
73 x[3] ^= x[0];
74 x[3] = @bitCast(Vector(4, u32), @shuffle(u8, @bitCast(Vector(16, u8), x[3]), undefined, rot16));
75
76 x[2] +%= x[3];
77 x[1] ^= x[2];
78
79 t1 = x[1];
80 x[1] <<= @splat(4, @as(u5, 12));
81 t1 >>= @splat(4, @as(u5, 20));
82 x[1] ^= t1;
83
84 x[0] +%= x[1];
85 x[3] ^= x[0];
86 x[0] = @shuffle(u32, x[0], undefined, Vector(4, i32){ 1, 2, 3, 0 });
87 x[3] = @bitCast(Vector(4, u32), @shuffle(u8, @bitCast(Vector(16, u8), x[3]), undefined, rot8));
88
89 x[2] +%= x[3];
90 x[3] = @shuffle(u32, x[3], undefined, Vector(4, i32){ 2, 3, 0, 1 });
91 x[1] ^= x[2];
92 x[2] = @shuffle(u32, x[2], undefined, Vector(4, i32){ 3, 0, 1, 2 });
93
94 t1 = x[1];
95 x[1] <<= @splat(4, @as(u5, 7));
96 t1 >>= @splat(4, @as(u5, 25));
97 x[1] ^= t1;
98 }
99 }
30100
31fn initContext(key: [8]u32, d: [4]u32) [16]u32 {101 inline fn hashToBytes(out: *[64]u8, x: BlockVec) void {
32 var ctx: [16]u32 = undefined;102 var i: usize = 0;
33 const c = "expand 32-byte k";103 while (i < 4) : (i += 1) {
34 const constant_le = comptime [_]u32{104 mem.writeIntLittle(u32, out[16 * i + 0 ..][0..4], x[i][0]);
35 mem.readIntLittle(u32, c[0..4]),105 mem.writeIntLittle(u32, out[16 * i + 4 ..][0..4], x[i][1]);
36 mem.readIntLittle(u32, c[4..8]),106 mem.writeIntLittle(u32, out[16 * i + 8 ..][0..4], x[i][2]);
37 mem.readIntLittle(u32, c[8..12]),107 mem.writeIntLittle(u32, out[16 * i + 12 ..][0..4], x[i][3]);
38 mem.readIntLittle(u32, c[12..16]),108 }
39 };109 }
40 mem.copy(u32, ctx[0..], constant_le[0..4]);
41 mem.copy(u32, ctx[4..12], key[0..8]);
42 mem.copy(u32, ctx[12..16], d[0..4]);
43110
44 return ctx;111 inline fn contextFeedback(x: *BlockVec, ctx: BlockVec) void {
45}112 x[0] +%= ctx[0];
113 x[1] +%= ctx[1];
114 x[2] +%= ctx[2];
115 x[3] +%= ctx[3];
116 }
46117
47// The chacha family of ciphers are based on the salsa family.118 fn chaCha20Internal(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
48inline fn chacha20Core(x: []u32, input: [16]u32) void {119 var ctx = initContext(key, counter);
49 for (x) |_, i|120 var x: BlockVec = undefined;
50 x[i] = input[i];121 var buf: [64]u8 = undefined;
51122 var i: usize = 0;
52 const rounds = comptime [_]QuarterRound{123 while (i + 64 <= in.len) : (i += 64) {
53 Rp(0, 4, 8, 12),124 chacha20Core(x[0..], ctx);
54 Rp(1, 5, 9, 13),125 contextFeedback(&x, ctx);
55 Rp(2, 6, 10, 14),126 hashToBytes(buf[0..], x);
56 Rp(3, 7, 11, 15),127
57 Rp(0, 5, 10, 15),128 var xout = out[i..];
58 Rp(1, 6, 11, 12),129 const xin = in[i..];
59 Rp(2, 7, 8, 13),130 var j: usize = 0;
60 Rp(3, 4, 9, 14),131 while (j < 64) : (j += 1) {
61 };132 xout[j] = xin[j];
133 }
134 j = 0;
135 while (j < 64) : (j += 1) {
136 xout[j] ^= buf[j];
137 }
138 ctx[3][0] += 1;
139 }
140 if (i < in.len) {
141 chacha20Core(x[0..], ctx);
142 contextFeedback(&x, ctx);
143 hashToBytes(buf[0..], x);
144
145 var xout = out[i..];
146 const xin = in[i..];
147 var j: usize = 0;
148 while (j < in.len % 64) : (j += 1) {
149 xout[j] = xin[j] ^ buf[j];
150 }
151 }
152 }
62153
63 comptime var j: usize = 0;154 fn hchacha20(input: [16]u8, key: [32]u8) [32]u8 {
64 inline while (j < 20) : (j += 2) {155 var c: [4]u32 = undefined;
65 // two-round cycles156 for (c) |_, i| {
66 inline for (rounds) |r| {157 c[i] = mem.readIntLittle(u32, input[4 * i ..][0..4]);
67 x[r.a] +%= x[r.b];
68 x[r.d] = std.math.rotl(u32, x[r.d] ^ x[r.a], @as(u32, 16));
69 x[r.c] +%= x[r.d];
70 x[r.b] = std.math.rotl(u32, x[r.b] ^ x[r.c], @as(u32, 12));
71 x[r.a] +%= x[r.b];
72 x[r.d] = std.math.rotl(u32, x[r.d] ^ x[r.a], @as(u32, 8));
73 x[r.c] +%= x[r.d];
74 x[r.b] = std.math.rotl(u32, x[r.b] ^ x[r.c], @as(u32, 7));
75 }158 }
159 const ctx = initContext(keyToWords(key), c);
160 var x: BlockVec = undefined;
161 chacha20Core(x[0..], ctx);
162 var out: [32]u8 = undefined;
163 mem.writeIntLittle(u32, out[0..4], x[0][0]);
164 mem.writeIntLittle(u32, out[4..8], x[0][1]);
165 mem.writeIntLittle(u32, out[8..12], x[0][2]);
166 mem.writeIntLittle(u32, out[12..16], x[0][3]);
167 mem.writeIntLittle(u32, out[16..20], x[3][0]);
168 mem.writeIntLittle(u32, out[20..24], x[3][1]);
169 mem.writeIntLittle(u32, out[24..28], x[3][2]);
170 mem.writeIntLittle(u32, out[28..32], x[3][3]);
171 return out;
76 }172 }
77}173};
78174
79fn hashToBytes(out: []u8, x: [16]u32) void {175// Non-vectorized implementation of the core function
80 for (x) |_, i| {176const ChaCha20NonVecImpl = struct {
81 mem.writeIntLittle(u32, out[4 * i ..][0..4], x[i]);177 const BlockVec = [16]u32;
178
179 fn initContext(key: [8]u32, d: [4]u32) BlockVec {
180 const c = "expand 32-byte k";
181 const constant_le = comptime [4]u32{
182 mem.readIntLittle(u32, c[0..4]),
183 mem.readIntLittle(u32, c[4..8]),
184 mem.readIntLittle(u32, c[8..12]),
185 mem.readIntLittle(u32, c[12..16]),
186 };
187 return BlockVec{
188 constant_le[0], constant_le[1], constant_le[2], constant_le[3],
189 key[0], key[1], key[2], key[3],
190 key[4], key[5], key[6], key[7],
191 d[0], d[1], d[2], d[3],
192 };
82 }193 }
83}
84194
85fn chaCha20_internal(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {195 const QuarterRound = struct {
86 var ctx = initContext(key, counter);196 a: usize,
87 var remaining: usize = if (in.len > out.len) in.len else out.len;197 b: usize,
88 var cursor: usize = 0;198 c: usize,
199 d: usize,
200 };
89201
90 while (true) {202 fn Rp(a: usize, b: usize, c: usize, d: usize) QuarterRound {
91 var x: [16]u32 = undefined;203 return QuarterRound{
92 var buf: [64]u8 = undefined;204 .a = a,
93 chacha20Core(x[0..], ctx);205 .b = b,
94 for (x) |_, i| {206 .c = c,
95 x[i] +%= ctx[i];207 .d = d,
208 };
209 }
210
211 inline fn chacha20Core(x: *BlockVec, input: BlockVec) void {
212 x.* = input;
213
214 const rounds = comptime [_]QuarterRound{
215 Rp(0, 4, 8, 12),
216 Rp(1, 5, 9, 13),
217 Rp(2, 6, 10, 14),
218 Rp(3, 7, 11, 15),
219 Rp(0, 5, 10, 15),
220 Rp(1, 6, 11, 12),
221 Rp(2, 7, 8, 13),
222 Rp(3, 4, 9, 14),
223 };
224
225 comptime var j: usize = 0;
226 inline while (j < 20) : (j += 2) {
227 inline for (rounds) |r| {
228 x[r.a] +%= x[r.b];
229 x[r.d] = std.math.rotl(u32, x[r.d] ^ x[r.a], @as(u32, 16));
230 x[r.c] +%= x[r.d];
231 x[r.b] = std.math.rotl(u32, x[r.b] ^ x[r.c], @as(u32, 12));
232 x[r.a] +%= x[r.b];
233 x[r.d] = std.math.rotl(u32, x[r.d] ^ x[r.a], @as(u32, 8));
234 x[r.c] +%= x[r.d];
235 x[r.b] = std.math.rotl(u32, x[r.b] ^ x[r.c], @as(u32, 7));
236 }
96 }237 }
97 hashToBytes(buf[0..], x);238 }
98 if (remaining < 64) {239
99 var i: usize = 0;240 inline fn hashToBytes(out: *[64]u8, x: BlockVec) void {
100 while (i < remaining) : (i += 1)241 var i: usize = 0;
101 out[cursor + i] = in[cursor + i] ^ buf[i];242 while (i < 4) : (i += 1) {
102 return;243 mem.writeIntLittle(u32, out[16 * i + 0 ..][0..4], x[i * 4 + 0]);
244 mem.writeIntLittle(u32, out[16 * i + 4 ..][0..4], x[i * 4 + 1]);
245 mem.writeIntLittle(u32, out[16 * i + 8 ..][0..4], x[i * 4 + 2]);
246 mem.writeIntLittle(u32, out[16 * i + 12 ..][0..4], x[i * 4 + 3]);
103 }247 }
248 }
104249
250 inline fn contextFeedback(x: *BlockVec, ctx: BlockVec) void {
105 var i: usize = 0;251 var i: usize = 0;
106 while (i < 64) : (i += 1)252 while (i < 16) : (i += 1) {
107 out[cursor + i] = in[cursor + i] ^ buf[i];253 x[i] +%= ctx[i];
254 }
255 }
108256
109 cursor += 64;257 fn chaCha20Internal(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
110 remaining -= 64;258 var ctx = initContext(key, counter);
259 var x: BlockVec = undefined;
260 var buf: [64]u8 = undefined;
261 var i: usize = 0;
262 while (i + 64 <= in.len) : (i += 64) {
263 chacha20Core(x[0..], ctx);
264 contextFeedback(&x, ctx);
265 hashToBytes(buf[0..], x);
266
267 var xout = out[i..];
268 const xin = in[i..];
269 var j: usize = 0;
270 while (j < 64) : (j += 1) {
271 xout[j] = xin[j];
272 }
273 j = 0;
274 while (j < 64) : (j += 1) {
275 xout[j] ^= buf[j];
276 }
277 ctx[12] += 1;
278 }
279 if (i < in.len) {
280 chacha20Core(x[0..], ctx);
281 contextFeedback(&x, ctx);
282 hashToBytes(buf[0..], x);
283
284 var xout = out[i..];
285 const xin = in[i..];
286 var j: usize = 0;
287 while (j < in.len % 64) : (j += 1) {
288 xout[j] = xin[j] ^ buf[j];
289 }
290 }
291 }
111292
112 ctx[12] += 1;293 fn hchacha20(input: [16]u8, key: [32]u8) [32]u8 {
294 var c: [4]u32 = undefined;
295 for (c) |_, i| {
296 c[i] = mem.readIntLittle(u32, input[4 * i ..][0..4]);
297 }
298 const ctx = initContext(keyToWords(key), c);
299 var x: BlockVec = undefined;
300 chacha20Core(x[0..], ctx);
301 var out: [32]u8 = undefined;
302 mem.writeIntLittle(u32, out[0..4], x[0]);
303 mem.writeIntLittle(u32, out[4..8], x[1]);
304 mem.writeIntLittle(u32, out[8..12], x[2]);
305 mem.writeIntLittle(u32, out[12..16], x[3]);
306 mem.writeIntLittle(u32, out[16..20], x[12]);
307 mem.writeIntLittle(u32, out[20..24], x[13]);
308 mem.writeIntLittle(u32, out[24..28], x[14]);
309 mem.writeIntLittle(u32, out[28..32], x[15]);
310 return out;
113 }311 }
114}312};
313
314const ChaCha20Impl = if (std.Target.current.cpu.arch == .x86_64) ChaCha20VecImpl else ChaCha20NonVecImpl;
115315
116fn keyToWords(key: [32]u8) [8]u32 {316fn keyToWords(key: [32]u8) [8]u32 {
117 var k: [8]u32 = undefined;317 var k: [8]u32 = undefined;
118 k[0] = mem.readIntLittle(u32, key[0..4]);318 var i: usize = 0;
119 k[1] = mem.readIntLittle(u32, key[4..8]);319 while (i < 8) : (i += 1) {
120 k[2] = mem.readIntLittle(u32, key[8..12]);320 k[i] = mem.readIntLittle(u32, key[i * 4 ..][0..4]);
121 k[3] = mem.readIntLittle(u32, key[12..16]);321 }
122 k[4] = mem.readIntLittle(u32, key[16..20]);
123 k[5] = mem.readIntLittle(u32, key[20..24]);
124 k[6] = mem.readIntLittle(u32, key[24..28]);
125 k[7] = mem.readIntLittle(u32, key[28..32]);
126
127 return k;322 return k;
128}323}
129324
...@@ -145,7 +340,7 @@ pub const ChaCha20IETF = struct {...@@ -145,7 +340,7 @@ pub const ChaCha20IETF = struct {
145 c[1] = mem.readIntLittle(u32, nonce[0..4]);340 c[1] = mem.readIntLittle(u32, nonce[0..4]);
146 c[2] = mem.readIntLittle(u32, nonce[4..8]);341 c[2] = mem.readIntLittle(u32, nonce[4..8]);
147 c[3] = mem.readIntLittle(u32, nonce[8..12]);342 c[3] = mem.readIntLittle(u32, nonce[8..12]);
148 chaCha20_internal(out, in, keyToWords(key), c);343 ChaCha20Impl.chaCha20Internal(out, in, keyToWords(key), c);
149 }344 }
150};345};
151346
...@@ -171,7 +366,7 @@ pub const ChaCha20With64BitNonce = struct {...@@ -171,7 +366,7 @@ pub const ChaCha20With64BitNonce = struct {
171366
172 // first partial big block367 // first partial big block
173 if (((@intCast(u64, maxInt(u32) - @truncate(u32, counter)) + 1) << 6) < in.len) {368 if (((@intCast(u64, maxInt(u32) - @truncate(u32, counter)) + 1) << 6) < in.len) {
174 chaCha20_internal(out[cursor..big_block], in[cursor..big_block], k, c);369 ChaCha20Impl.chaCha20Internal(out[cursor..big_block], in[cursor..big_block], k, c);
175 cursor = big_block - cursor;370 cursor = big_block - cursor;
176 c[1] += 1;371 c[1] += 1;
177 if (comptime @sizeOf(usize) > 4) {372 if (comptime @sizeOf(usize) > 4) {
...@@ -179,14 +374,14 @@ pub const ChaCha20With64BitNonce = struct {...@@ -179,14 +374,14 @@ pub const ChaCha20With64BitNonce = struct {
179 var remaining_blocks: u32 = @intCast(u32, (in.len / big_block));374 var remaining_blocks: u32 = @intCast(u32, (in.len / big_block));
180 var i: u32 = 0;375 var i: u32 = 0;
181 while (remaining_blocks > 0) : (remaining_blocks -= 1) {376 while (remaining_blocks > 0) : (remaining_blocks -= 1) {
182 chaCha20_internal(out[cursor .. cursor + big_block], in[cursor .. cursor + big_block], k, c);377 ChaCha20Impl.chaCha20Internal(out[cursor .. cursor + big_block], in[cursor .. cursor + big_block], k, c);
183 c[1] += 1; // upper 32-bit of counter, generic chaCha20_internal() doesn't know about this.378 c[1] += 1; // upper 32-bit of counter, generic chaCha20Internal() doesn't know about this.
184 cursor += big_block;379 cursor += big_block;
185 }380 }
186 }381 }
187 }382 }
188383
189 chaCha20_internal(out[cursor..], in[cursor..], k, c);384 ChaCha20Impl.chaCha20Internal(out[cursor..], in[cursor..], k, c);
190 }385 }
191};386};
192387
...@@ -533,33 +728,12 @@ fn chacha20poly1305Open(dst: []u8, ciphertextAndTag: []const u8, data: []const u...@@ -533,33 +728,12 @@ fn chacha20poly1305Open(dst: []u8, ciphertextAndTag: []const u8, data: []const u
533 return try chacha20poly1305OpenDetached(dst, ciphertextAndTag[0..ciphertextLen], ciphertextAndTag[ciphertextLen..][0..chacha20poly1305_tag_size], data, key, nonce);728 return try chacha20poly1305OpenDetached(dst, ciphertextAndTag[0..ciphertextLen], ciphertextAndTag[ciphertextLen..][0..chacha20poly1305_tag_size], data, key, nonce);
534}729}
535730
536fn hchacha20(input: [16]u8, key: [32]u8) [32]u8 {
537 var c: [4]u32 = undefined;
538 for (c) |_, i| {
539 c[i] = mem.readIntLittle(u32, input[4 * i ..][0..4]);
540 }
541 const ctx = initContext(keyToWords(key), c);
542 var x: [16]u32 = undefined;
543 chacha20Core(x[0..], ctx);
544 var out: [32]u8 = undefined;
545 mem.writeIntLittle(u32, out[0..4], x[0]);
546 mem.writeIntLittle(u32, out[4..8], x[1]);
547 mem.writeIntLittle(u32, out[8..12], x[2]);
548 mem.writeIntLittle(u32, out[12..16], x[3]);
549 mem.writeIntLittle(u32, out[16..20], x[12]);
550 mem.writeIntLittle(u32, out[20..24], x[13]);
551 mem.writeIntLittle(u32, out[24..28], x[14]);
552 mem.writeIntLittle(u32, out[28..32], x[15]);
553
554 return out;
555}
556
557fn extend(key: [32]u8, nonce: [24]u8) struct { key: [32]u8, nonce: [12]u8 } {731fn extend(key: [32]u8, nonce: [24]u8) struct { key: [32]u8, nonce: [12]u8 } {
558 var subnonce: [12]u8 = undefined;732 var subnonce: [12]u8 = undefined;
559 mem.set(u8, subnonce[0..4], 0);733 mem.set(u8, subnonce[0..4], 0);
560 mem.copy(u8, subnonce[4..], nonce[16..24]);734 mem.copy(u8, subnonce[4..], nonce[16..24]);
561 return .{735 return .{
562 .key = hchacha20(nonce[0..16].*, key),736 .key = ChaCha20Impl.hchacha20(nonce[0..16].*, key),
563 .nonce = subnonce,737 .nonce = subnonce,
564 };738 };
565}739}