authorgravatar for 124872+jedisct1@users.noreply.github.comFrank Denis <124872+jedisct1@users.noreply.github.com> 2025-11-03 17:09:00+01:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2025-11-03 17:09:00+01:00
logee4df4ad3edad160fb737a1935cd86bc2f9cfbbe
tree88a30b506f3d620a74adaf2b01dbb6f2c17d43c7
parentafdd04356cbf7821e1f992d1f3aedf7226a7e564
signaturebadge-check Signed by PGP key B5690EEEBB952194

crypto - threaded K12: separate context computation from thread spawning (#25793)

* threaded K12: separate context computation from thread spawning Compute all contexts and store them in a pre-allocated array, then spawn threads using the pre-computed contexts. This ensures each context is fully materialized in memory with the correct values before any thread tries to access it. * kt128: unroll the permutation rounds only twice This appears to deliver the best performance thanks to improved cache utilization, and it’s consistent with what we already do for SHA3.

1 files changed, 100 insertions(+), 89 deletions(-)

lib/std/crypto/kangarootwelve.zig+100-89
...@@ -230,58 +230,61 @@ fn keccakP1600timesN(comptime N: usize, states: *[5][5]@Vector(N, u64)) void {...@@ -230,58 +230,61 @@ fn keccakP1600timesN(comptime N: usize, states: *[5][5]@Vector(N, u64)) void {
230 break :blk offsets;230 break :blk offsets;
231 };231 };
232232
233 inline for (RC) |rc| {233 var round: usize = 0;
234 // θ (theta)234 while (round < 12) : (round += 2) {
235 var C: [5]@Vector(N, u64) = undefined;235 inline for (0..2) |i| {
236 inline for (0..5) |x| {236 // θ (theta)
237 C[x] = states[x][0] ^ states[x][1] ^ states[x][2] ^ states[x][3] ^ states[x][4];237 var C: [5]@Vector(N, u64) = undefined;
238 }238 inline for (0..5) |x| {
239 C[x] = states[x][0] ^ states[x][1] ^ states[x][2] ^ states[x][3] ^ states[x][4];
240 }
239241
240 var D: [5]@Vector(N, u64) = undefined;242 var D: [5]@Vector(N, u64) = undefined;
241 inline for (0..5) |x| {243 inline for (0..5) |x| {
242 D[x] = C[(x + 4) % 5] ^ rol64Vec(N, C[(x + 1) % 5], 1);244 D[x] = C[(x + 4) % 5] ^ rol64Vec(N, C[(x + 1) % 5], 1);
243 }245 }
244246
245 // Apply D to all lanes247 // Apply D to all lanes
246 inline for (0..5) |x| {248 inline for (0..5) |x| {
247 states[x][0] ^= D[x];249 states[x][0] ^= D[x];
248 states[x][1] ^= D[x];250 states[x][1] ^= D[x];
249 states[x][2] ^= D[x];251 states[x][2] ^= D[x];
250 states[x][3] ^= D[x];252 states[x][3] ^= D[x];
251 states[x][4] ^= D[x];253 states[x][4] ^= D[x];
252 }254 }
253255
254 // ρ (rho) and π (pi) - optimized with pre-computed offsets256 // ρ (rho) and π (pi) - optimized with pre-computed offsets
255 var current = states[1][0];257 var current = states[1][0];
256 var px: usize = 1;258 var px: usize = 1;
257 var py: usize = 0;259 var py: usize = 0;
258 inline for (rho_offsets) |rot| {260 inline for (rho_offsets) |rot| {
259 const next_y = (2 * px + 3 * py) % 5;261 const next_y = (2 * px + 3 * py) % 5;
260 const next = states[py][next_y];262 const next = states[py][next_y];
261 states[py][next_y] = rol64Vec(N, current, rot);263 states[py][next_y] = rol64Vec(N, current, rot);
262 current = next;264 current = next;
263 px = py;265 px = py;
264 py = next_y;266 py = next_y;
265 }267 }
266268
267 // χ (chi) - optimized with better register usage269 // χ (chi) - optimized with better register usage
268 inline for (0..5) |y| {270 inline for (0..5) |y| {
269 const t0 = states[0][y];271 const t0 = states[0][y];
270 const t1 = states[1][y];272 const t1 = states[1][y];
271 const t2 = states[2][y];273 const t2 = states[2][y];
272 const t3 = states[3][y];274 const t3 = states[3][y];
273 const t4 = states[4][y];275 const t4 = states[4][y];
274276
275 states[0][y] = t0 ^ (~t1 & t2);277 states[0][y] = t0 ^ (~t1 & t2);
276 states[1][y] = t1 ^ (~t2 & t3);278 states[1][y] = t1 ^ (~t2 & t3);
277 states[2][y] = t2 ^ (~t3 & t4);279 states[2][y] = t2 ^ (~t3 & t4);
278 states[3][y] = t3 ^ (~t4 & t0);280 states[3][y] = t3 ^ (~t4 & t0);
279 states[4][y] = t4 ^ (~t0 & t1);281 states[4][y] = t4 ^ (~t0 & t1);
280 }282 }
281283
282 // ι (iota)284 // ι (iota)
283 const rc_splat: @Vector(N, u64) = @splat(rc);285 const rc_splat: @Vector(N, u64) = @splat(RC[round + i]);
284 states[0][0] ^= rc_splat;286 states[0][0] ^= rc_splat;
287 }
285 }288 }
286}289}
287290
...@@ -323,46 +326,49 @@ fn keccakP(state: *[200]u8) void {...@@ -323,46 +326,49 @@ fn keccakP(state: *[200]u8) void {
323 }326 }
324327
325 // Apply 12 rounds328 // Apply 12 rounds
326 inline for (RC) |rc| {329 var round: usize = 0;
327 // θ330 while (round < 12) : (round += 2) {
328 var C: [5]u64 = undefined;331 inline for (0..2) |i| {
329 inline for (0..5) |x| {332 // θ
330 C[x] = lanes[x][0] ^ lanes[x][1] ^ lanes[x][2] ^ lanes[x][3] ^ lanes[x][4];333 var C: [5]u64 = undefined;
331 }334 inline for (0..5) |x| {
332 var D: [5]u64 = undefined;335 C[x] = lanes[x][0] ^ lanes[x][1] ^ lanes[x][2] ^ lanes[x][3] ^ lanes[x][4];
333 inline for (0..5) |x| {336 }
334 D[x] = C[(x + 4) % 5] ^ std.math.rotl(u64, C[(x + 1) % 5], 1);337 var D: [5]u64 = undefined;
335 }338 inline for (0..5) |x| {
336 inline for (0..5) |x| {339 D[x] = C[(x + 4) % 5] ^ std.math.rotl(u64, C[(x + 1) % 5], 1);
337 inline for (0..5) |y| {340 }
338 lanes[x][y] ^= D[x];341 inline for (0..5) |x| {
342 inline for (0..5) |y| {
343 lanes[x][y] ^= D[x];
344 }
339 }345 }
340 }
341346
342 // ρ and π347 // ρ and π
343 var current = lanes[1][0];348 var current = lanes[1][0];
344 var px: usize = 1;349 var px: usize = 1;
345 var py: usize = 0;350 var py: usize = 0;
346 inline for (0..24) |t| {351 inline for (0..24) |t| {
347 const temp = lanes[py][(2 * px + 3 * py) % 5];352 const temp = lanes[py][(2 * px + 3 * py) % 5];
348 const rot_amount = ((t + 1) * (t + 2) / 2) % 64;353 const rot_amount = ((t + 1) * (t + 2) / 2) % 64;
349 lanes[py][(2 * px + 3 * py) % 5] = std.math.rotl(u64, current, @as(u6, @intCast(rot_amount)));354 lanes[py][(2 * px + 3 * py) % 5] = std.math.rotl(u64, current, @as(u6, @intCast(rot_amount)));
350 current = temp;355 current = temp;
351 const temp_x = py;356 const temp_x = py;
352 py = (2 * px + 3 * py) % 5;357 py = (2 * px + 3 * py) % 5;
353 px = temp_x;358 px = temp_x;
354 }359 }
355360
356 // χ361 // χ
357 inline for (0..5) |y| {362 inline for (0..5) |y| {
358 const T = [5]u64{ lanes[0][y], lanes[1][y], lanes[2][y], lanes[3][y], lanes[4][y] };363 const T = [5]u64{ lanes[0][y], lanes[1][y], lanes[2][y], lanes[3][y], lanes[4][y] };
359 inline for (0..5) |x| {364 inline for (0..5) |x| {
360 lanes[x][y] = T[x] ^ (~T[(x + 1) % 5] & T[(x + 2) % 5]);365 lanes[x][y] = T[x] ^ (~T[(x + 1) % 5] & T[(x + 2) % 5]);
366 }
361 }367 }
362 }
363368
364 // ι369 // ι
365 lanes[0][0] ^= rc;370 lanes[0][0] ^= RC[round + i];
371 }
366 }372 }
367373
368 // Store lanes back to state374 // Store lanes back to state
...@@ -759,32 +765,37 @@ fn ktMultiThreaded(...@@ -759,32 +765,37 @@ fn ktMultiThreaded(
759 const all_scratch = try allocator.alloc(u8, thread_count * scratch_size);765 const all_scratch = try allocator.alloc(u8, thread_count * scratch_size);
760 defer allocator.free(all_scratch);766 defer allocator.free(all_scratch);
761767
762 var group: Io.Group = .init;768 const contexts = try allocator.alloc(LeafBatchContext, thread_count);
769 defer allocator.free(contexts);
770
763 var leaves_assigned: usize = 0;771 var leaves_assigned: usize = 0;
764 var thread_idx: usize = 0;772 var context_count: usize = 0;
765773
766 while (leaves_assigned < total_leaves) {774 while (leaves_assigned < total_leaves) {
767 const batch_count = @min(leaves_per_thread, total_leaves - leaves_assigned);775 const batch_count = @min(leaves_per_thread, total_leaves - leaves_assigned);
768 const batch_start = chunk_size + leaves_assigned * chunk_size;776 const batch_start = chunk_size + leaves_assigned * chunk_size;
769 const cvs_offset = leaves_assigned * cv_size;777 const cvs_offset = leaves_assigned * cv_size;
770778
771 const ctx = LeafBatchContext{779 contexts[context_count] = LeafBatchContext{
772 .output_cvs = cvs[cvs_offset .. cvs_offset + batch_count * cv_size],780 .output_cvs = cvs[cvs_offset .. cvs_offset + batch_count * cv_size],
773 .batch_start = batch_start,781 .batch_start = batch_start,
774 .batch_count = batch_count,782 .batch_count = batch_count,
775 .view = view,783 .view = view,
776 .scratch_buffer = all_scratch[thread_idx * scratch_size .. (thread_idx + 1) * scratch_size],784 .scratch_buffer = all_scratch[context_count * scratch_size .. (context_count + 1) * scratch_size],
777 .total_len = total_len,785 .total_len = total_len,
778 };786 };
779787
788 leaves_assigned += batch_count;
789 context_count += 1;
790 }
791
792 var group: Io.Group = .init;
793 for (contexts[0..context_count]) |ctx| {
780 group.async(io, struct {794 group.async(io, struct {
781 fn process(c: LeafBatchContext) void {795 fn process(c: LeafBatchContext) void {
782 processLeafBatch(Variant, c);796 processLeafBatch(Variant, c);
783 }797 }
784 }.process, .{ctx});798 }.process, .{ctx});
785
786 leaves_assigned += batch_count;
787 thread_idx += 1;
788 }799 }
789800
790 // Wait for all threads to complete801 // Wait for all threads to complete