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;
1010const assert = std.debug.assert;
1111const testing = std.testing;
1212const maxInt = std.math.maxInt;
13const Vector = std.meta.Vector;
1314const Poly1305 = std.crypto.onetimeauth.Poly1305;
1415
15const QuarterRound = struct {
16 a: usize,
17 b: usize,
18 c: usize,
19 d: usize,
20};
16// Vectorized implementation of the core function
17const ChaCha20VecImpl = struct {
18 const Lane = Vector(4, u32);
19 const BlockVec = [4]Lane;
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 {
23 return QuarterRound{
24 .a = a,
25 .b = b,
26 .c = c,
27 .d = d,
28 };
29}
37 inline fn chacha20Core(x: *BlockVec, input: BlockVec) void {
38 const rot8 = Vector(16, i32){ 3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14 };
39 const rot16 = Vector(16, i32){ 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13 };
40
41 x.* = input;
42
43 var r: usize = 0;
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 {
32 var ctx: [16]u32 = undefined;
33 const c = "expand 32-byte k";
34 const constant_le = comptime [_]u32{
35 mem.readIntLittle(u32, c[0..4]),
36 mem.readIntLittle(u32, c[4..8]),
37 mem.readIntLittle(u32, c[8..12]),
38 mem.readIntLittle(u32, c[12..16]),
39 };
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]);
101 inline fn hashToBytes(out: *[64]u8, x: BlockVec) void {
102 var i: usize = 0;
103 while (i < 4) : (i += 1) {
104 mem.writeIntLittle(u32, out[16 * i + 0 ..][0..4], x[i][0]);
105 mem.writeIntLittle(u32, out[16 * i + 4 ..][0..4], x[i][1]);
106 mem.writeIntLittle(u32, out[16 * i + 8 ..][0..4], x[i][2]);
107 mem.writeIntLittle(u32, out[16 * i + 12 ..][0..4], x[i][3]);
108 }
109 }
43110
44 return ctx;
45}
111 inline fn contextFeedback(x: *BlockVec, ctx: BlockVec) void {
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.
48inline fn chacha20Core(x: []u32, input: [16]u32) void {
49 for (x) |_, i|
50 x[i] = input[i];
51
52 const rounds = comptime [_]QuarterRound{
53 Rp(0, 4, 8, 12),
54 Rp(1, 5, 9, 13),
55 Rp(2, 6, 10, 14),
56 Rp(3, 7, 11, 15),
57 Rp(0, 5, 10, 15),
58 Rp(1, 6, 11, 12),
59 Rp(2, 7, 8, 13),
60 Rp(3, 4, 9, 14),
61 };
118 fn chaCha20Internal(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
119 var ctx = initContext(key, counter);
120 var x: BlockVec = undefined;
121 var buf: [64]u8 = undefined;
122 var i: usize = 0;
123 while (i + 64 <= in.len) : (i += 64) {
124 chacha20Core(x[0..], ctx);
125 contextFeedback(&x, ctx);
126 hashToBytes(buf[0..], x);
127
128 var xout = out[i..];
129 const xin = in[i..];
130 var j: usize = 0;
131 while (j < 64) : (j += 1) {
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;
64 inline while (j < 20) : (j += 2) {
65 // two-round cycles
66 inline for (rounds) |r| {
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));
154 fn hchacha20(input: [16]u8, key: [32]u8) [32]u8 {
155 var c: [4]u32 = undefined;
156 for (c) |_, i| {
157 c[i] = mem.readIntLittle(u32, input[4 * i ..][0..4]);
75158 }
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;
76172 }
77}
173};
78174
79fn hashToBytes(out: []u8, x: [16]u32) void {
80 for (x) |_, i| {
81 mem.writeIntLittle(u32, out[4 * i ..][0..4], x[i]);
175// Non-vectorized implementation of the core function
176const ChaCha20NonVecImpl = struct {
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 };
82193 }
83}
84194
85fn chaCha20_internal(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
86 var ctx = initContext(key, counter);
87 var remaining: usize = if (in.len > out.len) in.len else out.len;
88 var cursor: usize = 0;
195 const QuarterRound = struct {
196 a: usize,
197 b: usize,
198 c: usize,
199 d: usize,
200 };
89201
90 while (true) {
91 var x: [16]u32 = undefined;
92 var buf: [64]u8 = undefined;
93 chacha20Core(x[0..], ctx);
94 for (x) |_, i| {
95 x[i] +%= ctx[i];
202 fn Rp(a: usize, b: usize, c: usize, d: usize) QuarterRound {
203 return QuarterRound{
204 .a = a,
205 .b = b,
206 .c = c,
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 }
96237 }
97 hashToBytes(buf[0..], x);
98 if (remaining < 64) {
99 var i: usize = 0;
100 while (i < remaining) : (i += 1)
101 out[cursor + i] = in[cursor + i] ^ buf[i];
102 return;
238 }
239
240 inline fn hashToBytes(out: *[64]u8, x: BlockVec) void {
241 var i: usize = 0;
242 while (i < 4) : (i += 1) {
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]);
103247 }
248 }
104249
250 inline fn contextFeedback(x: *BlockVec, ctx: BlockVec) void {
105251 var i: usize = 0;
106 while (i < 64) : (i += 1)
107 out[cursor + i] = in[cursor + i] ^ buf[i];
252 while (i < 16) : (i += 1) {
253 x[i] +%= ctx[i];
254 }
255 }
108256
109 cursor += 64;
110 remaining -= 64;
257 fn chaCha20Internal(out: []u8, in: []const u8, key: [8]u32, counter: [4]u32) void {
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;
113311 }
114}
312};
313
314const ChaCha20Impl = if (std.Target.current.cpu.arch == .x86_64) ChaCha20VecImpl else ChaCha20NonVecImpl;
115315
116316fn keyToWords(key: [32]u8) [8]u32 {
117317 var k: [8]u32 = undefined;
118 k[0] = mem.readIntLittle(u32, key[0..4]);
119 k[1] = mem.readIntLittle(u32, key[4..8]);
120 k[2] = mem.readIntLittle(u32, key[8..12]);
121 k[3] = mem.readIntLittle(u32, key[12..16]);
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
318 var i: usize = 0;
319 while (i < 8) : (i += 1) {
320 k[i] = mem.readIntLittle(u32, key[i * 4 ..][0..4]);
321 }
127322 return k;
128323}
129324
......@@ -145,7 +340,7 @@ pub const ChaCha20IETF = struct {
145340 c[1] = mem.readIntLittle(u32, nonce[0..4]);
146341 c[2] = mem.readIntLittle(u32, nonce[4..8]);
147342 c[3] = mem.readIntLittle(u32, nonce[8..12]);
148 chaCha20_internal(out, in, keyToWords(key), c);
343 ChaCha20Impl.chaCha20Internal(out, in, keyToWords(key), c);
149344 }
150345};
151346
......@@ -171,7 +366,7 @@ pub const ChaCha20With64BitNonce = struct {
171366
172367 // first partial big block
173368 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);
175370 cursor = big_block - cursor;
176371 c[1] += 1;
177372 if (comptime @sizeOf(usize) > 4) {
......@@ -179,14 +374,14 @@ pub const ChaCha20With64BitNonce = struct {
179374 var remaining_blocks: u32 = @intCast(u32, (in.len / big_block));
180375 var i: u32 = 0;
181376 while (remaining_blocks > 0) : (remaining_blocks -= 1) {
182 chaCha20_internal(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.
377 ChaCha20Impl.chaCha20Internal(out[cursor .. cursor + big_block], in[cursor .. cursor + big_block], k, c);
378 c[1] += 1; // upper 32-bit of counter, generic chaCha20Internal() doesn't know about this.
184379 cursor += big_block;
185380 }
186381 }
187382 }
188383
189 chaCha20_internal(out[cursor..], in[cursor..], k, c);
384 ChaCha20Impl.chaCha20Internal(out[cursor..], in[cursor..], k, c);
190385 }
191386};
192387
......@@ -533,33 +728,12 @@ fn chacha20poly1305Open(dst: []u8, ciphertextAndTag: []const u8, data: []const u
533728 return try chacha20poly1305OpenDetached(dst, ciphertextAndTag[0..ciphertextLen], ciphertextAndTag[ciphertextLen..][0..chacha20poly1305_tag_size], data, key, nonce);
534729}
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
557731fn extend(key: [32]u8, nonce: [24]u8) struct { key: [32]u8, nonce: [12]u8 } {
558732 var subnonce: [12]u8 = undefined;
559733 mem.set(u8, subnonce[0..4], 0);
560734 mem.copy(u8, subnonce[4..], nonce[16..24]);
561735 return .{
562 .key = hchacha20(nonce[0..16].*, key),
736 .key = ChaCha20Impl.hchacha20(nonce[0..16].*, key),
563737 .nonce = subnonce,
564738 };
565739}