| ... | @@ -23,6 +23,12 @@ const max_simd_degree_or_2 = if (max_simd_degree > 2) max_simd_degree else 2; | ... | @@ -23,6 +23,12 @@ const max_simd_degree_or_2 = if (max_simd_degree > 2) max_simd_degree else 2; |
| 23 | /// Benchmarks generally show significant speedup starting at 3 MiB. | 23 | /// Benchmarks generally show significant speedup starting at 3 MiB. |
| 24 | const parallel_threshold = 3 * 1024 * 1024; | 24 | const parallel_threshold = 3 * 1024 * 1024; |
| 25 | | 25 | |
| | 26 | /// Largest subtree fold that runs as a single task. |
| | 27 | const max_fold_task_length = 2048; |
| | 28 | |
| | 29 | /// Smallest fold worth its own thread. |
| | 30 | const min_fold_task_length = 256; |
| | 31 | |
| 26 | const iv: [8]u32 = .{ | 32 | const iv: [8]u32 = .{ |
| 27 | 0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, | 33 | 0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, |
| 28 | 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19, | 34 | 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19, |
| ... | @@ -677,6 +683,7 @@ const ChunkBatch = struct { | ... | @@ -677,6 +683,7 @@ const ChunkBatch = struct { |
| 677 | input: []const u8, | 683 | input: []const u8, |
| 678 | start_chunk: usize, | 684 | start_chunk: usize, |
| 679 | end_chunk: usize, | 685 | end_chunk: usize, |
| | 686 | chunk_counter: u64, |
| 680 | cvs: [][8]u32, | 687 | cvs: [][8]u32, |
| 681 | key: [8]u32, | 688 | key: [8]u32, |
| 682 | flags: Flags, | 689 | flags: Flags, |
| ... | @@ -694,7 +701,7 @@ const ChunkBatch = struct { | ... | @@ -694,7 +701,7 @@ const ChunkBatch = struct { |
| 694 | const num_cvs = compressChunksParallel( | 701 | const num_cvs = compressChunksParallel( |
| 695 | ctx.input[offset..][0..batch_len], | 702 | ctx.input[offset..][0..batch_len], |
| 696 | ctx.key, | 703 | ctx.key, |
| 697 | chunk_idx, | 704 | ctx.chunk_counter + chunk_idx, |
| 698 | ctx.flags, | 705 | ctx.flags, |
| 699 | &cv_buffer, | 706 | &cv_buffer, |
| 700 | ); | 707 | ); |
| ... | @@ -807,6 +814,39 @@ fn buildMerkleTreeLayerParallel( | ... | @@ -807,6 +814,39 @@ fn buildMerkleTreeLayerParallel( |
| 807 | try group.await(io); | 814 | try group.await(io); |
| 808 | } | 815 | } |
| 809 | | 816 | |
| | 817 | // Largest power-of-two chunk count that fits and is aligned at chunk_counter. |
| | 818 | fn nextSubtreeChunks(remaining_chunks: usize, chunk_counter: u64) usize { |
| | 819 | var subtree_chunks: usize = @intCast(roundDownToPowerOf2(remaining_chunks)); |
| | 820 | while ((subtree_chunks - 1) & chunk_counter != 0) { |
| | 821 | subtree_chunks /= 2; |
| | 822 | } |
| | 823 | return subtree_chunks; |
| | 824 | } |
| | 825 | |
| | 826 | // Fold a subtree's chaining values (not the final pair!) in place. |
| | 827 | const SubtreeFold = struct { |
| | 828 | cvs: [][8]u32, |
| | 829 | key: [8]u32, |
| | 830 | flags: Flags, |
| | 831 | |
| | 832 | fn process(ctx: SubtreeFold) void { |
| | 833 | var level = ctx.cvs; |
| | 834 | |
| | 835 | while (level.len > 2) { |
| | 836 | const num_parents = level.len / 2; |
| | 837 | processParentBatchSIMD(ParentBatchContext{ |
| | 838 | .input_cvs = level, |
| | 839 | .output_cvs = level[0..num_parents], |
| | 840 | .start_idx = 0, |
| | 841 | .end_idx = num_parents, |
| | 842 | .key = ctx.key, |
| | 843 | .flags = ctx.flags, |
| | 844 | }); |
| | 845 | level = level[0..num_parents]; |
| | 846 | } |
| | 847 | } |
| | 848 | }; |
| | 849 | |
| 810 | fn parentOutput(parent_block: []const u8, key: [8]u32, flags: Flags) Output { | 850 | fn parentOutput(parent_block: []const u8, key: [8]u32, flags: Flags) Output { |
| 811 | var block: [Blake3.block_length]u8 = undefined; | 851 | var block: [Blake3.block_length]u8 = undefined; |
| 812 | @memcpy(&block, parent_block[0..Blake3.block_length]); | 852 | @memcpy(&block, parent_block[0..Blake3.block_length]); |
| ... | @@ -991,91 +1031,9 @@ pub const Blake3 = struct { | ... | @@ -991,91 +1031,9 @@ pub const Blake3 = struct { |
| 991 | } | 1031 | } |
| 992 | | 1032 | |
| 993 | pub fn hashParallel(b: []const u8, out: []u8, options: Options, allocator: Allocator, io: Io) error{ OutOfMemory, Canceled }!void { | 1033 | pub fn hashParallel(b: []const u8, out: []u8, options: Options, allocator: Allocator, io: Io) error{ OutOfMemory, Canceled }!void { |
| 994 | if (b.len < parallel_threshold) { | 1034 | var d = Blake3.init(options); |
| 995 | return hash(b, out, options); | 1035 | try d.updateParallel(b, allocator, io); |
| 996 | } | 1036 | d.final(out); |
| 997 | | | |
| 998 | const key_words = if (options.key) |key| loadKeyWords(key) else iv; | | |
| 999 | const flags: Flags = if (options.key != null) .{ .keyed_hash = true } else .{}; | | |
| 1000 | | | |
| 1001 | const num_full_chunks = b.len / chunk_length; | | |
| 1002 | const thread_count = Thread.getCpuCount() catch 1; | | |
| 1003 | if (thread_count <= 1 or num_full_chunks == 0) { | | |
| 1004 | return hash(b, out, options); | | |
| 1005 | } | | |
| 1006 | | | |
| 1007 | const remaining_bytes = b.len % chunk_length; | | |
| 1008 | const num_leaves = @divCeil(b.len, chunk_length); | | |
| 1009 | | | |
| 1010 | const cvs = try allocator.alloc([8]u32, num_leaves); | | |
| 1011 | defer allocator.free(cvs); | | |
| 1012 | | | |
| 1013 | // Process chunks in parallel | | |
| 1014 | const num_workers = thread_count; | | |
| 1015 | const chunks_per_worker = (num_full_chunks + num_workers - 1) / num_workers; | | |
| 1016 | var group: Io.Group = .init; | | |
| 1017 | defer group.cancel(io); | | |
| 1018 | | | |
| 1019 | for (0..num_workers) |worker_id| { | | |
| 1020 | const start_chunk = worker_id * chunks_per_worker; | | |
| 1021 | if (start_chunk >= num_full_chunks) break; | | |
| 1022 | | | |
| 1023 | group.async(io, ChunkBatch.process, .{ChunkBatch{ | | |
| 1024 | .input = b, | | |
| 1025 | .start_chunk = start_chunk, | | |
| 1026 | .end_chunk = @min(start_chunk + chunks_per_worker, num_full_chunks), | | |
| 1027 | .cvs = cvs, | | |
| 1028 | .key = key_words, | | |
| 1029 | .flags = flags, | | |
| 1030 | }}); | | |
| 1031 | } | | |
| 1032 | try group.await(io); | | |
| 1033 | | | |
| 1034 | if (remaining_bytes > 0) { | | |
| 1035 | var chunk_state = ChunkState.init(key_words, flags); | | |
| 1036 | chunk_state.chunk_counter = num_full_chunks; | | |
| 1037 | chunk_state.update(b[num_full_chunks * chunk_length ..]); | | |
| 1038 | const output = chunk_state.output(); | | |
| 1039 | cvs[num_full_chunks] = output.chainingValue(); | | |
| 1040 | } | | |
| 1041 | | | |
| 1042 | // Build Merkle tree in parallel layers using ping-pong buffers | | |
| 1043 | const max_intermediate_size = @divCeil(num_leaves, 2); | | |
| 1044 | const buffer0 = try allocator.alloc([8]u32, max_intermediate_size); | | |
| 1045 | defer allocator.free(buffer0); | | |
| 1046 | const buffer1 = try allocator.alloc([8]u32, max_intermediate_size); | | |
| 1047 | defer allocator.free(buffer1); | | |
| 1048 | | | |
| 1049 | var current_level = cvs; | | |
| 1050 | var next_level_buf = buffer0; | | |
| 1051 | var toggle = false; | | |
| 1052 | | | |
| 1053 | while (current_level.len > 8) { | | |
| 1054 | const num_parents = current_level.len / 2; | | |
| 1055 | const has_odd = current_level.len % 2 == 1; | | |
| 1056 | const next_level_size = num_parents + @intFromBool(has_odd); | | |
| 1057 | | | |
| 1058 | try buildMerkleTreeLayerParallel( | | |
| 1059 | current_level[0 .. num_parents * 2], | | |
| 1060 | next_level_buf[0..num_parents], | | |
| 1061 | key_words, | | |
| 1062 | flags, | | |
| 1063 | io, | | |
| 1064 | ); | | |
| 1065 | | | |
| 1066 | if (has_odd) { | | |
| 1067 | next_level_buf[num_parents] = current_level[current_level.len - 1]; | | |
| 1068 | } | | |
| 1069 | | | |
| 1070 | current_level = next_level_buf[0..next_level_size]; | | |
| 1071 | next_level_buf = if (toggle) buffer0 else buffer1; | | |
| 1072 | toggle = !toggle; | | |
| 1073 | } | | |
| 1074 | | | |
| 1075 | // Finalize remaining small tree sequentially | | |
| 1076 | var hasher = init_internal(key_words, flags); | | |
| 1077 | for (current_level, 0..) |cv, i| hasher.pushCv(cv, i); | | |
| 1078 | hasher.final(out); | | |
| 1079 | } | 1037 | } |
| 1080 | | 1038 | |
| 1081 | fn init_internal(key: [8]u32, flags: Flags) Blake3 { | 1039 | fn init_internal(key: [8]u32, flags: Flags) Blake3 { |
| ... | @@ -1159,6 +1117,150 @@ pub const Blake3 = struct { | ... | @@ -1159,6 +1117,150 @@ pub const Blake3 = struct { |
| 1159 | } | 1117 | } |
| 1160 | } | 1118 | } |
| 1161 | | 1119 | |
| | 1120 | /// Add input to the hash state, hashing using multiple threads. This can be called any number of times. |
| | 1121 | pub fn updateParallel(self: *Blake3, input: []const u8, allocator: Allocator, io: Io) error{ OutOfMemory, Canceled }!void { |
| | 1122 | if (input.len < parallel_threshold) { |
| | 1123 | return self.update(input); |
| | 1124 | } |
| | 1125 | |
| | 1126 | const thread_count = Thread.getCpuCount() catch 1; |
| | 1127 | if (thread_count <= 1) { |
| | 1128 | return self.update(input); |
| | 1129 | } |
| | 1130 | |
| | 1131 | var inp = input; |
| | 1132 | |
| | 1133 | if (self.chunk.len() > 0) { |
| | 1134 | const take = chunk_length - self.chunk.len(); |
| | 1135 | self.chunk.update(inp[0..take]); |
| | 1136 | inp = inp[take..]; |
| | 1137 | const output = self.chunk.output(); |
| | 1138 | const chunk_cv = output.chainingValue(); |
| | 1139 | self.pushCv(chunk_cv, self.chunk.chunk_counter); |
| | 1140 | self.chunk.reset(self.key, self.chunk.chunk_counter + 1); |
| | 1141 | } |
| | 1142 | |
| | 1143 | const base_counter = self.chunk.chunk_counter; |
| | 1144 | const full_chunks = inp.len / chunk_length; |
| | 1145 | const keeps_trailing_chunk = inp.len % chunk_length == 0 and (base_counter + full_chunks) % 2 != 0; |
| | 1146 | const num_chunks = full_chunks - @intFromBool(keeps_trailing_chunk); |
| | 1147 | |
| | 1148 | const cvs = try allocator.alloc([8]u32, num_chunks); |
| | 1149 | defer allocator.free(cvs); |
| | 1150 | |
| | 1151 | const tree_buffer_len = if (num_chunks >= 2 * max_fold_task_length) num_chunks else 0; |
| | 1152 | const buffer0 = try allocator.alloc([8]u32, tree_buffer_len / 2); |
| | 1153 | defer allocator.free(buffer0); |
| | 1154 | const buffer1 = try allocator.alloc([8]u32, tree_buffer_len / 4); |
| | 1155 | defer allocator.free(buffer1); |
| | 1156 | |
| | 1157 | const chunks_per_worker = @divCeil(num_chunks, thread_count); |
| | 1158 | var group: Io.Group = .init; |
| | 1159 | defer group.cancel(io); |
| | 1160 | |
| | 1161 | for (0..thread_count) |worker_id| { |
| | 1162 | const start_chunk = worker_id * chunks_per_worker; |
| | 1163 | if (start_chunk >= num_chunks) break; |
| | 1164 | |
| | 1165 | group.async(io, ChunkBatch.process, .{ChunkBatch{ |
| | 1166 | .input = inp, |
| | 1167 | .start_chunk = start_chunk, |
| | 1168 | .end_chunk = @min(start_chunk + chunks_per_worker, num_chunks), |
| | 1169 | .chunk_counter = base_counter, |
| | 1170 | .cvs = cvs, |
| | 1171 | .key = self.key, |
| | 1172 | .flags = self.chunk.flags, |
| | 1173 | }}); |
| | 1174 | } |
| | 1175 | try group.await(io); |
| | 1176 | |
| | 1177 | var level: [][8]u32 = cvs; |
| | 1178 | var level_start: usize = 0; |
| | 1179 | var chunks_per_cv: usize = 1; |
| | 1180 | var next_level_buf = buffer0; |
| | 1181 | var spare_level_buf = buffer1; |
| | 1182 | |
| | 1183 | while (true) { |
| | 1184 | const level_end = level_start + level.len * chunks_per_cv; |
| | 1185 | var run_start: ?usize = null; |
| | 1186 | var run_end: usize = undefined; |
| | 1187 | var subtree_start = level_start; |
| | 1188 | |
| | 1189 | while (subtree_start < level_end) { |
| | 1190 | const subtree_chunks = nextSubtreeChunks(num_chunks - subtree_start, base_counter + subtree_start); |
| | 1191 | const subtree_cvs = level[(subtree_start - level_start) / chunks_per_cv ..][0 .. subtree_chunks / chunks_per_cv]; |
| | 1192 | |
| | 1193 | if (subtree_cvs.len > max_fold_task_length) { |
| | 1194 | if (run_start == null) { |
| | 1195 | run_start = subtree_start; |
| | 1196 | } else { |
| | 1197 | std.debug.assert(run_end == subtree_start); |
| | 1198 | } |
| | 1199 | run_end = subtree_start + subtree_chunks; |
| | 1200 | } else if (chunks_per_cv > 1) { |
| | 1201 | @memcpy(cvs[subtree_start..][0..subtree_cvs.len], subtree_cvs); |
| | 1202 | } |
| | 1203 | subtree_start += subtree_chunks; |
| | 1204 | } |
| | 1205 | const start = run_start orelse break; |
| | 1206 | |
| | 1207 | level = level[(start - level_start) / chunks_per_cv ..][0 .. (run_end - start) / chunks_per_cv]; |
| | 1208 | level_start = start; |
| | 1209 | |
| | 1210 | const num_parents = level.len / 2; |
| | 1211 | try buildMerkleTreeLayerParallel( |
| | 1212 | level, |
| | 1213 | next_level_buf[0..num_parents], |
| | 1214 | self.key, |
| | 1215 | self.chunk.flags, |
| | 1216 | io, |
| | 1217 | ); |
| | 1218 | level = next_level_buf[0..num_parents]; |
| | 1219 | chunks_per_cv *= 2; |
| | 1220 | mem.swap([][8]u32, &next_level_buf, &spare_level_buf); |
| | 1221 | } |
| | 1222 | |
| | 1223 | var fold_group: Io.Group = .init; |
| | 1224 | defer fold_group.cancel(io); |
| | 1225 | |
| | 1226 | var chunk_idx: usize = 0; |
| | 1227 | while (chunk_idx < num_chunks) { |
| | 1228 | const subtree_chunks = nextSubtreeChunks(num_chunks - chunk_idx, base_counter + chunk_idx); |
| | 1229 | const fold: SubtreeFold = .{ |
| | 1230 | .cvs = cvs[chunk_idx..][0..@min(subtree_chunks, max_fold_task_length)], |
| | 1231 | .key = self.key, |
| | 1232 | .flags = self.chunk.flags, |
| | 1233 | }; |
| | 1234 | |
| | 1235 | if (fold.cvs.len >= min_fold_task_length) { |
| | 1236 | fold_group.async(io, SubtreeFold.process, .{fold}); |
| | 1237 | } else { |
| | 1238 | SubtreeFold.process(fold); |
| | 1239 | } |
| | 1240 | chunk_idx += subtree_chunks; |
| | 1241 | } |
| | 1242 | try fold_group.await(io); |
| | 1243 | |
| | 1244 | chunk_idx = 0; |
| | 1245 | while (chunk_idx < num_chunks) { |
| | 1246 | const counter = base_counter + chunk_idx; |
| | 1247 | const subtree_chunks = nextSubtreeChunks(num_chunks - chunk_idx, counter); |
| | 1248 | |
| | 1249 | self.pushCv(cvs[chunk_idx], counter); |
| | 1250 | if (subtree_chunks > 1) { |
| | 1251 | self.pushCv(cvs[chunk_idx + 1], counter + subtree_chunks / 2); |
| | 1252 | } |
| | 1253 | chunk_idx += subtree_chunks; |
| | 1254 | } |
| | 1255 | |
| | 1256 | self.chunk.reset(self.key, base_counter + num_chunks); |
| | 1257 | const tail = inp[num_chunks * chunk_length ..]; |
| | 1258 | if (tail.len > 0) { |
| | 1259 | self.chunk.update(tail); |
| | 1260 | self.mergeCvStack(self.chunk.chunk_counter); |
| | 1261 | } |
| | 1262 | } |
| | 1263 | |
| 1162 | /// Finalize the hash and write any number of output bytes. | 1264 | /// Finalize the hash and write any number of output bytes. |
| 1163 | pub fn final(self: *const Blake3, out: []u8) void { | 1265 | pub fn final(self: *const Blake3, out: []u8) void { |
| 1164 | self.finalizeSeek(0, out); | 1266 | self.finalizeSeek(0, out); |
| ... | @@ -1487,3 +1589,24 @@ test "BLAKE3 parallel with partial trailing chunk" { | ... | @@ -1487,3 +1589,24 @@ test "BLAKE3 parallel with partial trailing chunk" { |
| 1487 | try std.testing.expectEqualSlices(u8, &expected, &actual); | 1589 | try std.testing.expectEqualSlices(u8, &expected, &actual); |
| 1488 | } | 1590 | } |
| 1489 | } | 1591 | } |
| | 1592 | |
| | 1593 | test "BLAKE3 incremental parallel update" { |
| | 1594 | const allocator = std.testing.allocator; |
| | 1595 | const io = std.testing.io; |
| | 1596 | |
| | 1597 | const input = try allocator.alloc(u8, 16385 * 1024 + 1); |
| | 1598 | defer allocator.free(input); |
| | 1599 | for (input, 0..) |*byte, i| byte.* = @truncate(i); |
| | 1600 | |
| | 1601 | var expected: [32]u8 = undefined; |
| | 1602 | Blake3.hash(input, &expected, .{}); |
| | 1603 | |
| | 1604 | const prefix_len = 2 * 1024 * 1024 + 1; |
| | 1605 | var hasher = Blake3.init(.{}); |
| | 1606 | hasher.update(input[0..prefix_len]); |
| | 1607 | try hasher.updateParallel(input[prefix_len..], allocator, io); |
| | 1608 | var actual: [32]u8 = undefined; |
| | 1609 | hasher.final(&actual); |
| | 1610 | |
| | 1611 | try std.testing.expectEqualSlices(u8, &expected, &actual); |
| | 1612 | } |
Frank Denis <124872+jedisct1@users.noreply.github.com> 2026-08-13 03:36:08+02:00