| ... | ... | @@ -23,6 +23,12 @@ const max_simd_degree_or_2 = if (max_simd_degree > 2) max_simd_degree else 2; |
| 23 | 23 | /// Benchmarks generally show significant speedup starting at 3 MiB. |
| 24 | 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 | 32 | const iv: [8]u32 = .{ |
| 27 | 33 | 0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, |
| 28 | 34 | 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19, |
| ... | ... | @@ -677,6 +683,7 @@ const ChunkBatch = struct { |
| 677 | 683 | input: []const u8, |
| 678 | 684 | start_chunk: usize, |
| 679 | 685 | end_chunk: usize, |
| 686 | chunk_counter: u64, |
| 680 | 687 | cvs: [][8]u32, |
| 681 | 688 | key: [8]u32, |
| 682 | 689 | flags: Flags, |
| ... | ... | @@ -694,7 +701,7 @@ const ChunkBatch = struct { |
| 694 | 701 | const num_cvs = compressChunksParallel( |
| 695 | 702 | ctx.input[offset..][0..batch_len], |
| 696 | 703 | ctx.key, |
| 697 | | chunk_idx, |
| 704 | ctx.chunk_counter + chunk_idx, |
| 698 | 705 | ctx.flags, |
| 699 | 706 | &cv_buffer, |
| 700 | 707 | ); |
| ... | ... | @@ -807,6 +814,39 @@ fn buildMerkleTreeLayerParallel( |
| 807 | 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 | 850 | fn parentOutput(parent_block: []const u8, key: [8]u32, flags: Flags) Output { |
| 811 | 851 | var block: [Blake3.block_length]u8 = undefined; |
| 812 | 852 | @memcpy(&block, parent_block[0..Blake3.block_length]); |
| ... | ... | @@ -991,91 +1031,9 @@ pub const Blake3 = struct { |
| 991 | 1031 | } |
| 992 | 1032 | |
| 993 | 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) { |
| 995 | | return hash(b, out, options); |
| 996 | | } |
| 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); |
| 1034 | var d = Blake3.init(options); |
| 1035 | try d.updateParallel(b, allocator, io); |
| 1036 | d.final(out); |
| 1079 | 1037 | } |
| 1080 | 1038 | |
| 1081 | 1039 | fn init_internal(key: [8]u32, flags: Flags) Blake3 { |
| ... | ... | @@ -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 | 1264 | /// Finalize the hash and write any number of output bytes. |
| 1163 | 1265 | pub fn final(self: *const Blake3, out: []u8) void { |
| 1164 | 1266 | self.finalizeSeek(0, out); |
| ... | ... | @@ -1487,3 +1589,24 @@ test "BLAKE3 parallel with partial trailing chunk" { |
| 1487 | 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