authorgravatar for 124872+jedisct1@users.noreply.github.comFrank Denis <124872+jedisct1@users.noreply.github.com> 2020-10-23 16:18:35+02:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-10-25 21:13:14-04:00
log72064eba233fd3fd8e3dc0ab6f5fc73c6bdf1922
treed4f479fb0ee636c75f23810e751190a6050eae2e
parent1b4ab749cffeff146b21fb0f6d30a5b8d2186765

std/crypto: vectorize BLAKE3

Gives a ~40% speedup on x86_64. However, the generic code remains faster on aarch64. This is still processing only one block at a time for now. I'm pretty confident that processing more blocks per round will eventually give a substantial performance improvement on all platforms with vector units.

1 files changed, 159 insertions(+), 76 deletions(-)

lib/std/crypto/blake3.zig+159-76
...@@ -11,6 +11,7 @@ const fmt = std.fmt;...@@ -11,6 +11,7 @@ const fmt = std.fmt;
11const math = std.math;11const math = std.math;
12const mem = std.mem;12const mem = std.mem;
13const testing = std.testing;13const testing = std.testing;
14const Vector = std.meta.Vector;
1415
15const ChunkIterator = struct {16const ChunkIterator = struct {
16 slice: []u8,17 slice: []u8,
...@@ -61,87 +62,173 @@ const KEYED_HASH: u8 = 1 << 4;...@@ -61,87 +62,173 @@ const KEYED_HASH: u8 = 1 << 4;
61const DERIVE_KEY_CONTEXT: u8 = 1 << 5;62const DERIVE_KEY_CONTEXT: u8 = 1 << 5;
62const DERIVE_KEY_MATERIAL: u8 = 1 << 6;63const DERIVE_KEY_MATERIAL: u8 = 1 << 6;
6364
64// The mixing function, G, which mixes either a column or a diagonal.65const CompressVectorized = struct {
65fn g(state: *[16]u32, a: usize, b: usize, c: usize, d: usize, mx: u32, my: u32) void {66 const Lane = Vector(4, u32);
66 _ = @addWithOverflow(u32, state[a], state[b], &state[a]);67 const Rows = [4]Lane;
67 _ = @addWithOverflow(u32, state[a], mx, &state[a]);
68 state[d] = math.rotr(u32, state[d] ^ state[a], 16);
69 _ = @addWithOverflow(u32, state[c], state[d], &state[c]);
70 state[b] = math.rotr(u32, state[b] ^ state[c], 12);
71 _ = @addWithOverflow(u32, state[a], state[b], &state[a]);
72 _ = @addWithOverflow(u32, state[a], my, &state[a]);
73 state[d] = math.rotr(u32, state[d] ^ state[a], 8);
74 _ = @addWithOverflow(u32, state[c], state[d], &state[c]);
75 state[b] = math.rotr(u32, state[b] ^ state[c], 7);
76}
7768
78fn round(state: *[16]u32, msg: [16]u32, schedule: [16]u8) void {69 inline fn rot(x: Lane, comptime n: u5) Lane {
79 // Mix the columns.70 return (x >> @splat(4, @as(u5, n))) | (x << @splat(4, @as(u5, 1 +% ~n)));
80 g(state, 0, 4, 8, 12, msg[schedule[0]], msg[schedule[1]]);71 }
81 g(state, 1, 5, 9, 13, msg[schedule[2]], msg[schedule[3]]);
82 g(state, 2, 6, 10, 14, msg[schedule[4]], msg[schedule[5]]);
83 g(state, 3, 7, 11, 15, msg[schedule[6]], msg[schedule[7]]);
84
85 // Mix the diagonals.
86 g(state, 0, 5, 10, 15, msg[schedule[8]], msg[schedule[9]]);
87 g(state, 1, 6, 11, 12, msg[schedule[10]], msg[schedule[11]]);
88 g(state, 2, 7, 8, 13, msg[schedule[12]], msg[schedule[13]]);
89 g(state, 3, 4, 9, 14, msg[schedule[14]], msg[schedule[15]]);
90}
9172
92fn compress(73 inline fn g(comptime even: bool, rows: *Rows, m: Lane) void {
93 chaining_value: [8]u32,74 rows[0] +%= rows[1] +% m;
94 block_words: [16]u32,75 rows[3] ^= rows[0];
95 block_len: u32,76 rows[3] = rot(rows[3], if (even) 8 else 16);
96 counter: u64,77 rows[2] +%= rows[3];
97 flags: u8,78 rows[1] ^= rows[2];
98) [16]u32 {79 rows[1] = rot(rows[1], if (even) 7 else 12);
99 var state = [16]u32{
100 chaining_value[0],
101 chaining_value[1],
102 chaining_value[2],
103 chaining_value[3],
104 chaining_value[4],
105 chaining_value[5],
106 chaining_value[6],
107 chaining_value[7],
108 IV[0],
109 IV[1],
110 IV[2],
111 IV[3],
112 @truncate(u32, counter),
113 @truncate(u32, counter >> 32),
114 block_len,
115 flags,
116 };
117 for (MSG_SCHEDULE) |schedule| {
118 round(&state, block_words, schedule);
119 }80 }
120 for (chaining_value) |_, i| {81
121 state[i] ^= state[i + 8];82 inline fn diagonalize(rows: *Rows) void {
122 state[i + 8] ^= chaining_value[i];83 rows[0] = @shuffle(u32, rows[0], undefined, [_]i32{ 3, 0, 1, 2 });
84 rows[3] = @shuffle(u32, rows[3], undefined, [_]i32{ 2, 3, 0, 1 });
85 rows[2] = @shuffle(u32, rows[2], undefined, [_]i32{ 1, 2, 3, 0 });
123 }86 }
124 return state;87
125}88 inline fn undiagonalize(rows: *Rows) void {
89 rows[0] = @shuffle(u32, rows[0], undefined, [_]i32{ 1, 2, 3, 0 });
90 rows[3] = @shuffle(u32, rows[3], undefined, [_]i32{ 2, 3, 0, 1 });
91 rows[2] = @shuffle(u32, rows[2], undefined, [_]i32{ 3, 0, 1, 2 });
92 }
93
94 fn compress(
95 chaining_value: [8]u32,
96 block_words: [16]u32,
97 block_len: u32,
98 counter: u64,
99 flags: u8,
100 ) [16]u32 {
101 const md = Lane{ @truncate(u32, counter), @truncate(u32, counter >> 32), block_len, @as(u32, flags) };
102 var rows = Rows{ chaining_value[0..4].*, chaining_value[4..8].*, IV[0..4].*, md };
103
104 var m = Rows{ block_words[0..4].*, block_words[4..8].*, block_words[8..12].*, block_words[12..16].* };
105 var t0 = @shuffle(u32, m[0], m[1], [_]i32{ 0, 2, (-1 - 0), (-1 - 2) });
106 g(false, &rows, t0);
107 var t1 = @shuffle(u32, m[0], m[1], [_]i32{ 1, 3, (-1 - 1), (-1 - 3) });
108 g(true, &rows, t1);
109 diagonalize(&rows);
110 var t2 = @shuffle(u32, m[2], m[3], [_]i32{ 0, 2, (-1 - 0), (-1 - 2) });
111 t2 = @shuffle(u32, t2, undefined, [_]i32{ 3, 0, 1, 2 });
112 g(false, &rows, t2);
113 var t3 = @shuffle(u32, m[2], m[3], [_]i32{ 1, 3, (-1 - 1), (-1 - 3) });
114 t3 = @shuffle(u32, t3, undefined, [_]i32{ 3, 0, 1, 2 });
115 g(true, &rows, t3);
116 undiagonalize(&rows);
117 m = Rows{ t0, t1, t2, t3 };
118
119 var i: usize = 0;
120 while (i < 6) : (i += 1) {
121 t0 = @shuffle(u32, m[0], m[1], [_]i32{ 2, 1, (-1 - 1), (-1 - 3) });
122 t0 = @shuffle(u32, t0, undefined, [_]i32{ 1, 2, 3, 0 });
123 g(false, &rows, t0);
124 t1 = @shuffle(u32, m[2], m[3], [_]i32{ 2, 2, (-1 - 3), (-1 - 3) });
125 var tt = @shuffle(u32, m[0], undefined, [_]i32{ 3, 3, 0, 0 });
126 t1 = @shuffle(u32, tt, t1, [_]i32{ 0, (-1 - 1), 2, (-1 - 3) });
127 g(true, &rows, t1);
128 diagonalize(&rows);
129 t2 = @shuffle(u32, m[3], m[1], [_]i32{ 0, 1, (-1 - 0), (-1 - 1) });
130 tt = @shuffle(u32, t2, m[2], [_]i32{ 0, 1, 2, (-1 - 3) });
131 t2 = @shuffle(u32, tt, undefined, [_]i32{ 0, 2, 3, 1 });
132 g(false, &rows, t2);
133 t3 = @shuffle(u32, m[1], m[3], [_]i32{ 2, (-1 - 2), 3, (-1 - 3) });
134 tt = @shuffle(u32, m[2], t3, [_]i32{ 0, (-1 - 0), 1, (-1 - 1) });
135 t3 = @shuffle(u32, tt, undefined, [_]i32{ 2, 3, 1, 0 });
136 g(true, &rows, t3);
137 undiagonalize(&rows);
138 m = Rows{ t0, t1, t2, t3 };
139 }
140
141 rows[0] ^= rows[2];
142 rows[1] ^= rows[3];
143 rows[2] ^= Vector(4, u32){ chaining_value[0], chaining_value[1], chaining_value[2], chaining_value[3] };
144 rows[3] ^= Vector(4, u32){ chaining_value[4], chaining_value[5], chaining_value[6], chaining_value[7] };
145
146 return @bitCast([16]u32, rows);
147 }
148};
149
150const CompressGeneric = struct {
151 fn g(state: *[16]u32, comptime a: usize, comptime b: usize, comptime c: usize, comptime d: usize, mx: u32, my: u32) void {
152 state[a] +%= state[b] +% mx;
153 state[d] = math.rotr(u32, state[d] ^ state[a], 16);
154 state[c] +%= state[d];
155 state[b] = math.rotr(u32, state[b] ^ state[c], 12);
156 state[a] +%= state[b] +% my;
157 state[d] = math.rotr(u32, state[d] ^ state[a], 8);
158 state[c] +%= state[d];
159 state[b] = math.rotr(u32, state[b] ^ state[c], 7);
160 }
161
162 fn round(state: *[16]u32, msg: [16]u32, schedule: [16]u8) void {
163 // Mix the columns.
164 g(state, 0, 4, 8, 12, msg[schedule[0]], msg[schedule[1]]);
165 g(state, 1, 5, 9, 13, msg[schedule[2]], msg[schedule[3]]);
166 g(state, 2, 6, 10, 14, msg[schedule[4]], msg[schedule[5]]);
167 g(state, 3, 7, 11, 15, msg[schedule[6]], msg[schedule[7]]);
168
169 // Mix the diagonals.
170 g(state, 0, 5, 10, 15, msg[schedule[8]], msg[schedule[9]]);
171 g(state, 1, 6, 11, 12, msg[schedule[10]], msg[schedule[11]]);
172 g(state, 2, 7, 8, 13, msg[schedule[12]], msg[schedule[13]]);
173 g(state, 3, 4, 9, 14, msg[schedule[14]], msg[schedule[15]]);
174 }
175
176 fn compress(
177 chaining_value: [8]u32,
178 block_words: [16]u32,
179 block_len: u32,
180 counter: u64,
181 flags: u8,
182 ) [16]u32 {
183 var state = [16]u32{
184 chaining_value[0],
185 chaining_value[1],
186 chaining_value[2],
187 chaining_value[3],
188 chaining_value[4],
189 chaining_value[5],
190 chaining_value[6],
191 chaining_value[7],
192 IV[0],
193 IV[1],
194 IV[2],
195 IV[3],
196 @truncate(u32, counter),
197 @truncate(u32, counter >> 32),
198 block_len,
199 flags,
200 };
201 for (MSG_SCHEDULE) |schedule| {
202 round(&state, block_words, schedule);
203 }
204 for (chaining_value) |_, i| {
205 state[i] ^= state[i + 8];
206 state[i + 8] ^= chaining_value[i];
207 }
208 return state;
209 }
210};
211
212const compress = if (std.Target.current.cpu.arch == .x86_64) CompressVectorized.compress else CompressGeneric.compress;
126213
127fn first8Words(words: [16]u32) [8]u32 {214fn first8Words(words: [16]u32) [8]u32 {
128 return @ptrCast(*const [8]u32, &words).*;215 return @ptrCast(*const [8]u32, &words).*;
129}216}
130217
131fn wordsFromLittleEndianBytes(words: []u32, bytes: []const u8) void {218fn wordsFromLittleEndianBytes(comptime count: usize, bytes: [count * 4]u8) [count]u32 {
132 var byte_slice = bytes;219 var words: [count]u32 = undefined;
133 for (words) |*word| {220 for (words) |*word, i| {
134 word.* = mem.readIntSliceLittle(u32, byte_slice);221 word.* = mem.readIntSliceLittle(u32, bytes[4 * i ..]);
135 byte_slice = byte_slice[4..];
136 }222 }
223 return words;
137}224}
138225
139// Each chunk or parent node can produce either an 8-word chaining value or, by226// Each chunk or parent node can produce either an 8-word chaining value or, by
140// setting the ROOT flag, any number of final output bytes. The Output struct227// setting the ROOT flag, any number of final output bytes. The Output struct
141// captures the state just prior to choosing between those two possibilities.228// captures the state just prior to choosing between those two possibilities.
142const Output = struct {229const Output = struct {
143 input_chaining_value: [8]u32,230 input_chaining_value: [8]u32 align(16),
144 block_words: [16]u32,231 block_words: [16]u32 align(16),
145 block_len: u32,232 block_len: u32,
146 counter: u64,233 counter: u64,
147 flags: u8,234 flags: u8,
...@@ -181,9 +268,9 @@ const Output = struct {...@@ -181,9 +268,9 @@ const Output = struct {
181};268};
182269
183const ChunkState = struct {270const ChunkState = struct {
184 chaining_value: [8]u32,271 chaining_value: [8]u32 align(16),
185 chunk_counter: u64,272 chunk_counter: u64,
186 block: [BLOCK_LEN]u8 = [_]u8{0} ** BLOCK_LEN,273 block: [BLOCK_LEN]u8 align(16) = [_]u8{0} ** BLOCK_LEN,
187 block_len: u8 = 0,274 block_len: u8 = 0,
188 blocks_compressed: u8 = 0,275 blocks_compressed: u8 = 0,
189 flags: u8,276 flags: u8,
...@@ -218,8 +305,7 @@ const ChunkState = struct {...@@ -218,8 +305,7 @@ const ChunkState = struct {
218 // If the block buffer is full, compress it and clear it. More305 // If the block buffer is full, compress it and clear it. More
219 // input is coming, so this compression is not CHUNK_END.306 // input is coming, so this compression is not CHUNK_END.
220 if (self.block_len == BLOCK_LEN) {307 if (self.block_len == BLOCK_LEN) {
221 var block_words: [16]u32 = undefined;308 const block_words = wordsFromLittleEndianBytes(16, self.block);
222 wordsFromLittleEndianBytes(block_words[0..], self.block[0..]);
223 self.chaining_value = first8Words(compress(309 self.chaining_value = first8Words(compress(
224 self.chaining_value,310 self.chaining_value,
225 block_words,311 block_words,
...@@ -238,8 +324,7 @@ const ChunkState = struct {...@@ -238,8 +324,7 @@ const ChunkState = struct {
238 }324 }
239325
240 fn output(self: *const ChunkState) Output {326 fn output(self: *const ChunkState) Output {
241 var block_words: [16]u32 = undefined;327 const block_words = wordsFromLittleEndianBytes(16, self.block);
242 wordsFromLittleEndianBytes(block_words[0..], self.block[0..]);
243 return Output{328 return Output{
244 .input_chaining_value = self.chaining_value,329 .input_chaining_value = self.chaining_value,
245 .block_words = block_words,330 .block_words = block_words,
...@@ -256,7 +341,7 @@ fn parentOutput(...@@ -256,7 +341,7 @@ fn parentOutput(
256 key: [8]u32,341 key: [8]u32,
257 flags: u8,342 flags: u8,
258) Output {343) Output {
259 var block_words: [16]u32 = undefined;344 var block_words: [16]u32 align(16) = undefined;
260 mem.copy(u32, block_words[0..8], left_child_cv[0..]);345 mem.copy(u32, block_words[0..8], left_child_cv[0..]);
261 mem.copy(u32, block_words[8..], right_child_cv[0..]);346 mem.copy(u32, block_words[8..], right_child_cv[0..]);
262 return Output{347 return Output{
...@@ -303,8 +388,7 @@ pub const Blake3 = struct {...@@ -303,8 +388,7 @@ pub const Blake3 = struct {
303 /// Construct a new `Blake3` for the hash function, with an optional key388 /// Construct a new `Blake3` for the hash function, with an optional key
304 pub fn init(options: Options) Blake3 {389 pub fn init(options: Options) Blake3 {
305 if (options.key) |key| {390 if (options.key) |key| {
306 var key_words: [8]u32 = undefined;391 const key_words = wordsFromLittleEndianBytes(8, key);
307 wordsFromLittleEndianBytes(key_words[0..], key[0..]);
308 return Blake3.init_internal(key_words, KEYED_HASH);392 return Blake3.init_internal(key_words, KEYED_HASH);
309 } else {393 } else {
310 return Blake3.init_internal(IV, 0);394 return Blake3.init_internal(IV, 0);
...@@ -318,8 +402,7 @@ pub const Blake3 = struct {...@@ -318,8 +402,7 @@ pub const Blake3 = struct {
318 context_hasher.update(context);402 context_hasher.update(context);
319 var context_key: [KEY_LEN]u8 = undefined;403 var context_key: [KEY_LEN]u8 = undefined;
320 context_hasher.final(context_key[0..]);404 context_hasher.final(context_key[0..]);
321 var context_key_words: [8]u32 = undefined;405 const context_key_words = wordsFromLittleEndianBytes(8, context_key);
322 wordsFromLittleEndianBytes(context_key_words[0..], context_key[0..]);
323 return Blake3.init_internal(context_key_words, DERIVE_KEY_MATERIAL);406 return Blake3.init_internal(context_key_words, DERIVE_KEY_MATERIAL);
324 }407 }
325408