lib/simd/src/sort/vq.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const constants = @import("constants.zig");
3 const heap = @import("heap.zig");
4 const key = @import("key.zig");
5 const network = @import("network.zig");
6 const partition = @import("partition.zig");
7 const pivot = @import("pivot.zig");
8 const random = @import("random.zig");
9
10 pub const Error = error{InvalidK};
11 pub const maximum_levels: usize = 50;
12
13 pub fn sort(
14 comptime Key: type,
15 comptime direction: key.Direction,
16 values: []Key,
17 ) void {
18 var scratch = network.Scratch{};
19 var state = initialState(values.len);
20 const valid = values.len - partition.nanToBack(Key, values);
21 sortRange(Key, direction, values[0..valid], &scratch, &state, maximum_levels);
22 }
23
24 pub fn partialSort(
25 comptime Key: type,
26 comptime direction: key.Direction,
27 values: []Key,
28 k: usize,
29 ) Error!void {
30 if (k > values.len) return error.InvalidK;
31 var scratch = network.Scratch{};
32 var state = initialState(values.len);
33 const valid = values.len - partition.nanToBack(Key, values);
34 const prefix = @min(k, valid);
35 if (prefix < valid) {
36 selectRange(Key, direction, values[0..valid], prefix, &scratch, &state, maximum_levels);
37 }
38 sortRange(Key, direction, values[0..prefix], &scratch, &state, maximum_levels);
39 }
40
41 pub fn select(
42 comptime Key: type,
43 comptime direction: key.Direction,
44 values: []Key,
45 k: usize,
46 ) Error!void {
47 if (k >= values.len) return error.InvalidK;
48 var scratch = network.Scratch{};
49 var state = initialState(values.len);
50 const valid = values.len - partition.nanToBack(Key, values);
51 if (k >= valid) return;
52 selectRange(Key, direction, values[0..valid], k, &scratch, &state, maximum_levels);
53 }
54
55 pub fn haveFloat16() bool {
56 return true;
57 }
58
59 pub fn haveFloat64() bool {
60 return true;
61 }
62
63 fn initialState(count: usize) random.State {
64 const size: u64 = @intCast(count);
65 return random.State.init(
66 0x243f_6a88_85a3_08d3 ^ size,
67 0x1319_8a2e_0370_7344 ^ (size *% 0x9e37_79b9_7f4a_7c15),
68 );
69 }
70
71 fn baseSort(
72 comptime Key: type,
73 comptime direction: key.Direction,
74 values: []Key,
75 scratch: *network.Scratch,
76 ) void {
77 network.sort(Key, direction, values, scratch) catch unreachable;
78 }
79
80 fn sortRange(
81 comptime Key: type,
82 comptime direction: key.Direction,
83 values: []Key,
84 scratch: *network.Scratch,
85 state: *random.State,
86 levels: usize,
87 ) void {
88 const capacity = comptime constants.baseCaseKeyCapacity(Key);
89 var active = values;
90 var remaining = levels;
91 while (active.len > capacity) {
92 if (remaining == 0) {
93 heap.sort(Key, direction, active);
94 return;
95 }
96 const selected = pivot.choose(Key, direction, active, state);
97 const bounds = partition.around(Key, direction, active, selected);
98 std.debug.assert(bounds.before < bounds.after);
99 remaining -= 1;
100 const left = active[0..bounds.before];
101 const right = active[bounds.after..];
102 if (left.len < right.len) {
103 if (left.len > 1) sortRange(Key, direction, left, scratch, state, remaining);
104 active = right;
105 } else {
106 if (right.len > 1) sortRange(Key, direction, right, scratch, state, remaining);
107 active = left;
108 }
109 }
110 baseSort(Key, direction, active, scratch);
111 }
112
113 fn selectRange(
114 comptime Key: type,
115 comptime direction: key.Direction,
116 values: []Key,
117 k: usize,
118 scratch: *network.Scratch,
119 state: *random.State,
120 levels: usize,
121 ) void {
122 std.debug.assert(k < values.len);
123 const capacity = comptime constants.baseCaseKeyCapacity(Key);
124 var active = values;
125 var selected_index = k;
126 var remaining = levels;
127 while (active.len > capacity) {
128 if (remaining == 0) {
129 heap.sort(Key, direction, active);
130 return;
131 }
132 const selected = pivot.choose(Key, direction, active, state);
133 const bounds = partition.around(Key, direction, active, selected);
134 std.debug.assert(bounds.before < bounds.after);
135 remaining -= 1;
136 if (selected_index < bounds.before) {
137 active = active[0..bounds.before];
138 } else if (selected_index >= bounds.after) {
139 active = active[bounds.after..];
140 selected_index -= bounds.after;
141 } else {
142 return;
143 }
144 }
145 baseSort(Key, direction, active, scratch);
146 }
147
148 fn Order(comptime Key: type, comptime direction: key.Direction) type {
149 return struct {
150 fn lessThan(_: void, a: Key, b: Key) bool {
151 return key.before(Key, direction, a, b);
152 }
153 };
154 }
155
156 fn fillU32(values: []u32) void {
157 var state = random.State.init(0x736f_6d65_7073_6575, 0x646f_7261_6e64_6f6d);
158 for (values) |*value| value.* = @truncate(state.next());
159 }
160
161 fn verifyIntegerSort(comptime Key: type) !void {
162 const count = comptime constants.baseCaseKeyCapacity(Key) + 17;
163 const UInt = @Int(.unsigned, @bitSizeOf(Key));
164 var source: [count]Key = undefined;
165 var state = random.State.init(0xa409_3822_299f_31d0, 0x082e_fa98_ec4e_6c89);
166 for (&source) |*value| {
167 const bits: UInt = @truncate(state.next());
168 value.* = @bitCast(bits);
169 }
170 inline for (.{ key.Direction.ascending, key.Direction.descending }) |direction| {
171 var values = source;
172 var expected = source;
173 std.mem.sort(Key, &expected, {}, Order(Key, direction).lessThan);
174 sort(Key, direction, &values);
175 try std.testing.expectEqualSlices(Key, &expected, &values);
176 }
177 }
178
179 test "Highway VQSort public full sort crosses the base-case boundary" {
180 var ascending: [1025]u32 = undefined;
181 fillU32(&ascending);
182 var expected_ascending = ascending;
183 std.mem.sort(u32, &expected_ascending, {}, Order(u32, .ascending).lessThan);
184 sort(u32, .ascending, &ascending);
185 try std.testing.expectEqualSlices(u32, &expected_ascending, &ascending);
186
187 var descending = expected_ascending;
188 var expected_descending = descending;
189 std.mem.sort(u32, &expected_descending, {}, Order(u32, .descending).lessThan);
190 sort(u32, .descending, &descending);
191 try std.testing.expectEqualSlices(u32, &expected_descending, &descending);
192 }
193
194 test "Highway VQSort public full sort covers every integer width" {
195 try verifyIntegerSort(u16);
196 try verifyIntegerSort(i16);
197 try verifyIntegerSort(u32);
198 try verifyIntegerSort(i32);
199 try verifyIntegerSort(u64);
200 try verifyIntegerSort(i64);
201 }
202
203 test "Highway VQSort public partial sort and select preserve their boundaries" {
204 var source: [1537]u32 = undefined;
205 fillU32(&source);
206 var expected = source;
207 std.mem.sort(u32, &expected, {}, Order(u32, .ascending).lessThan);
208
209 var partial = source;
210 try partialSort(u32, .ascending, &partial, 613);
211 try std.testing.expectEqualSlices(u32, expected[0..613], partial[0..613]);
212 for (partial[613..]) |value| try std.testing.expect(value >= partial[612]);
213
214 var selected = source;
215 try select(u32, .ascending, &selected, 613);
216 try std.testing.expectEqual(expected[613], selected[613]);
217 for (selected[0..613]) |value| try std.testing.expect(value <= selected[613]);
218 for (selected[614..]) |value| try std.testing.expect(value >= selected[613]);
219
220 var all = source;
221 try partialSort(u32, .ascending, &all, all.len);
222 try std.testing.expectEqualSlices(u32, &expected, &all);
223 }
224
225 test "Highway VQSort public recursion terminates equal and two-value inputs" {
226 var equal = @as([4097]u32, @splat(17));
227 sort(u32, .ascending, &equal);
228 try partialSort(u32, .descending, &equal, 2048);
229 try select(u32, .ascending, &equal, 3000);
230 for (equal) |value| try std.testing.expectEqual(@as(u32, 17), value);
231
232 var two: [4097]u32 = undefined;
233 for (&two, 0..) |*value, index| value.* = @intCast(index & 1);
234 sort(u32, .descending, &two);
235 for (two[1..], 1..) |value, index| try std.testing.expect(two[index - 1] >= value);
236 }
237
238 test "Highway VQSort public recursion reaches the heap fallback" {
239 var values: [513]u32 = undefined;
240 fillU32(&values);
241 var expected = values;
242 std.mem.sort(u32, &expected, {}, Order(u32, .ascending).lessThan);
243 var scratch = network.Scratch{};
244 var state = initialState(values.len);
245 sortRange(u32, .ascending, &values, &scratch, &state, 0);
246 try std.testing.expectEqualSlices(u32, &expected, &values);
247
248 fillU32(&values);
249 std.mem.sort(u32, &expected, {}, Order(u32, .descending).lessThan);
250 state = initialState(values.len);
251 selectRange(u32, .descending, &values, 271, &scratch, &state, 0);
252 try std.testing.expectEqual(expected[271], values[271]);
253 }
254
255 test "Highway VQSort public key records and uint128 retain key order" {
256 var records: [513]key.K32V32 = undefined;
257 for (&records, 0..) |*value, index| {
258 value.* = key.K32V32.init(@intCast((index * 97) % 31), @intCast(index));
259 }
260 const original = records;
261 sort(key.K32V32, .ascending, &records);
262 for (records[1..], 1..) |value, index| try std.testing.expect(records[index - 1].key <= value.key);
263 var expected_records = original;
264 var actual_records = records;
265 std.mem.sort(key.K32V32, &expected_records, {}, struct {
266 fn lessThan(_: void, a: key.K32V32, b: key.K32V32) bool {
267 return key.networkBefore(key.K32V32, .ascending, a, b);
268 }
269 }.lessThan);
270 std.mem.sort(key.K32V32, &actual_records, {}, struct {
271 fn lessThan(_: void, a: key.K32V32, b: key.K32V32) bool {
272 return key.networkBefore(key.K32V32, .ascending, a, b);
273 }
274 }.lessThan);
275 try std.testing.expectEqualSlices(key.K32V32, &expected_records, &actual_records);
276
277 var partial_records = original;
278 try partialSort(key.K32V32, .ascending, &partial_records, 211);
279 for (partial_records[1..211], 1..) |value, index| {
280 try std.testing.expect(partial_records[index - 1].key <= value.key);
281 }
282 for (partial_records[211..]) |value| try std.testing.expect(value.key >= partial_records[210].key);
283 var selected_records = original;
284 try select(key.K32V32, .descending, &selected_records, 211);
285 for (selected_records[0..211]) |value| try std.testing.expect(value.key >= selected_records[211].key);
286 for (selected_records[212..]) |value| try std.testing.expect(value.key <= selected_records[211].key);
287
288 var wide: [257]u128 = undefined;
289 for (&wide, 0..) |*value, index| {
290 value.* = (@as(u128, @intCast(index * 131 + 7)) << 64) | @as(u128, @intCast(257 - index));
291 }
292 var expected_wide = wide;
293 std.mem.sort(u128, &expected_wide, {}, Order(u128, .descending).lessThan);
294 sort(u128, .descending, &wide);
295 try std.testing.expectEqualSlices(u128, &expected_wide, &wide);
296
297 var records64: [129]key.K64V64 = undefined;
298 for (&records64, 0..) |*value, index| {
299 value.* = key.K64V64.init(@intCast((index * 193) % 43), @intCast(index));
300 }
301 sort(key.K64V64, .descending, &records64);
302 for (records64[1..], 1..) |value, index| {
303 try std.testing.expect(records64[index - 1].key >= value.key);
304 }
305 }
306
307 fn verifyFloatNaN(comptime Float: type, comptime direction: key.Direction) !void {
308 var source: [521]Float = undefined;
309 for (&source, 0..) |*value, index| {
310 const magnitude: Float = @floatFromInt(index % 127);
311 value.* = if (index & 1 == 0) magnitude else -magnitude;
312 }
313 source[7] = std.math.nan(Float);
314 source[91] = std.math.nan(Float);
315 source[173] = std.math.inf(Float);
316 source[281] = std.math.nan(Float);
317 source[407] = -std.math.inf(Float);
318 source[503] = std.math.nan(Float);
319
320 var expected_valid: [517]Float = undefined;
321 var write: usize = 0;
322 for (source) |value| {
323 if (!std.math.isNan(value)) {
324 expected_valid[write] = value;
325 write += 1;
326 }
327 }
328 try std.testing.expectEqual(expected_valid.len, write);
329 std.mem.sort(Float, &expected_valid, {}, Order(Float, direction).lessThan);
330
331 var full = source;
332 sort(Float, direction, &full);
333 try std.testing.expectEqualSlices(Float, &expected_valid, full[0..expected_valid.len]);
334 for (full[expected_valid.len..]) |value| try std.testing.expect(std.math.isNan(value));
335
336 var partial = source;
337 try partialSort(Float, direction, &partial, 311);
338 try std.testing.expectEqualSlices(Float, expected_valid[0..311], partial[0..311]);
339 for (partial[expected_valid.len..]) |value| try std.testing.expect(std.math.isNan(value));
340
341 var partial_all = source;
342 try partialSort(Float, direction, &partial_all, partial_all.len);
343 try std.testing.expectEqualSlices(Float, &expected_valid, partial_all[0..expected_valid.len]);
344 for (partial_all[expected_valid.len..]) |value| try std.testing.expect(std.math.isNan(value));
345
346 var selected = source;
347 try select(Float, direction, &selected, 311);
348 try std.testing.expectEqual(expected_valid[311], selected[311]);
349 for (selected[expected_valid.len..]) |value| try std.testing.expect(std.math.isNan(value));
350
351 var nan_selected = source;
352 try select(Float, direction, &nan_selected, 519);
353 try std.testing.expect(std.math.isNan(nan_selected[519]));
354 var infinity_count: usize = 0;
355 for (nan_selected[0..expected_valid.len]) |value| {
356 if (std.math.isInf(value)) infinity_count += 1;
357 }
358 try std.testing.expectEqual(@as(usize, 2), infinity_count);
359 }
360
361 test "Highway VQSort public float operations move NaN behind real infinities" {
362 inline for (.{ f16, f32, f64 }) |Float| {
363 try verifyFloatNaN(Float, .ascending);
364 try verifyFloatNaN(Float, .descending);
365 }
366 }
367
368 test "Highway VQSort public k contracts reject before mutation" {
369 var values = [_]f32{ 3, std.math.nan(f32), 1 };
370 const expected_bits = [_]u32{ @bitCast(values[0]), @bitCast(values[1]), @bitCast(values[2]) };
371 try std.testing.expectError(error.InvalidK, partialSort(f32, .ascending, &values, 4));
372 try std.testing.expectError(error.InvalidK, select(f32, .ascending, &values, 3));
373 for (values, expected_bits) |value, bits| try std.testing.expectEqual(bits, @as(u32, @bitCast(value)));
374 try partialSort(f32, .ascending, &values, 0);
375 try std.testing.expect(haveFloat16());
376 try std.testing.expect(haveFloat64());
377 }
378
379 fn oracleFold(digest: *u64, value: anytype) void {
380 digest.* = (digest.* ^ @as(u64, @intCast(value))) *% 0x0000_0100_0000_01b3;
381 }
382
383 fn oracleMix(input: u64) u64 {
384 var value = input;
385 value = (value ^ (value >> 30)) *% 0xbf58_476d_1ce4_e5b9;
386 value = (value ^ (value >> 27)) *% 0x94d0_49bb_1331_11eb;
387 return value ^ (value >> 31);
388 }
389
390 fn foldIntegerOracle(digest: *u64) void {
391 const count = 1021;
392 var u16_values: [count]u16 = undefined;
393 var i16_values: [count]i16 = undefined;
394 var u32_values: [count]u32 = undefined;
395 var i32_values: [count]i32 = undefined;
396 var u64_values: [count]u64 = undefined;
397 var i64_values: [count]i64 = undefined;
398 for (0..count) |index| {
399 const permutation = (index * 73 + 19) % count;
400 u16_values[index] = @intCast(permutation);
401 i16_values[index] = @intCast(@as(i32, @intCast(permutation)) - 510);
402 u32_values[index] = @truncate(oracleMix(index + 11));
403 i32_values[index] = @bitCast(@as(u32, @truncate(oracleMix(index + 23))));
404 u64_values[index] = oracleMix(index + 37);
405 i64_values[index] = @bitCast(oracleMix(index + 41));
406 }
407 sort(u16, .ascending, &u16_values);
408 sort(i16, .descending, &i16_values);
409 sort(u32, .ascending, &u32_values);
410 sort(i32, .descending, &i32_values);
411 sort(u64, .ascending, &u64_values);
412 sort(i64, .descending, &i64_values);
413 for (0..count) |index| {
414 oracleFold(digest, u16_values[index]);
415 oracleFold(digest, @as(u16, @bitCast(i16_values[index])));
416 oracleFold(digest, u32_values[index]);
417 oracleFold(digest, @as(u32, @bitCast(i32_values[index])));
418 oracleFold(digest, u64_values[index]);
419 oracleFold(digest, @as(u64, @bitCast(i64_values[index])));
420 }
421 }
422
423 fn foldFloatOracle(digest: *u64) void {
424 const count = 1021;
425 var f32_values: [count]f32 = undefined;
426 var f64_values: [count]f64 = undefined;
427 for (0..count) |index| {
428 const value: i32 = @as(i32, @intCast((index * 97 + 3) % count)) - 510;
429 f32_values[index] = @floatFromInt(value);
430 f64_values[index] = @floatFromInt(value);
431 }
432 f32_values[17] = std.math.inf(f32);
433 f32_values[93] = std.math.nan(f32);
434 f32_values[511] = std.math.nan(f32);
435 f32_values[900] = -std.math.inf(f32);
436 f64_values[29] = std.math.inf(f64);
437 f64_values[113] = std.math.nan(f64);
438 f64_values[617] = std.math.nan(f64);
439 f64_values[901] = -std.math.inf(f64);
440 sort(f32, .ascending, &f32_values);
441 sort(f64, .descending, &f64_values);
442 for (0..count) |index| {
443 oracleFold(digest, @as(u32, @bitCast(f32_values[index])));
444 oracleFold(digest, @as(u64, @bitCast(f64_values[index])));
445 }
446 }
447
448 fn foldWideOracle(digest: *u64) !void {
449 const count = 1021;
450 var kv32: [count]key.K32V32 = undefined;
451 for (&kv32, 0..) |*value, index| {
452 value.* = .{
453 .value = @truncate(oracleMix(index + 51)),
454 .key = @intCast((index * 101 + 7) % count),
455 };
456 }
457 sort(key.K32V32, .ascending, &kv32);
458 for (kv32) |value| {
459 oracleFold(digest, value.value);
460 oracleFold(digest, value.key);
461 }
462 try std.testing.expectEqual(@as(u64, 6_939_292_972_457_705_748), digest.*);
463
464 const wide_count = 257;
465 var u128_values: [wide_count]u128 = undefined;
466 var kv64: [wide_count]key.K64V64 = undefined;
467 for (0..wide_count) |index| {
468 u128_values[index] = @as(u128, oracleMix(index + 61)) |
469 (@as(u128, oracleMix(index + 71)) << 64);
470 kv64[index] = .{
471 .value = oracleMix(index + 81),
472 .key = @intCast((index * 193 + 11) % wide_count),
473 };
474 }
475 sort(u128, .descending, &u128_values);
476 sort(key.K64V64, .descending, &kv64);
477 for (0..wide_count) |index| {
478 oracleFold(digest, @as(u64, @truncate(u128_values[index])));
479 oracleFold(digest, @as(u64, @truncate(u128_values[index] >> 64)));
480 oracleFold(digest, kv64[index].value);
481 oracleFold(digest, kv64[index].key);
482 }
483 }
484
485 fn foldPartialSelectOracle(digest: *u64) !void {
486 const count = 1021;
487 var partial: [count]u32 = undefined;
488 for (&partial, 0..) |*value, index| value.* = @truncate(oracleMix(index + 91));
489 try partialSort(u32, .ascending, &partial, 317);
490 for (partial[0..317]) |value| oracleFold(digest, value);
491 try std.testing.expectEqual(@as(u64, 15_502_069_084_655_779_296), digest.*);
492
493 var selected: [count]u64 = undefined;
494 for (&selected, 0..) |*value, index| value.* = oracleMix(index + 101);
495 try select(u64, .descending, &selected, 503);
496 oracleFold(digest, selected[503]);
497 var violations: u64 = 0;
498 for (selected[0..503]) |value| violations += @intFromBool(value < selected[503]);
499 for (selected[504..]) |value| violations += @intFromBool(value > selected[503]);
500 oracleFold(digest, violations);
501 try std.testing.expectEqual(@as(u64, 13_890_455_305_228_908_481), digest.*);
502
503 const nan_count = 521;
504 var nan_values: [nan_count]f32 = undefined;
505 for (&nan_values, 0..) |*value, index| value.* = @floatFromInt((index * 67 + 5) % 509);
506 nan_values[13] = std.math.nan(f32);
507 nan_values[107] = std.math.inf(f32);
508 nan_values[211] = std.math.nan(f32);
509 nan_values[401] = -std.math.inf(f32);
510 nan_values[503] = std.math.nan(f32);
511 var nan_partial = nan_values;
512 try partialSort(f32, .descending, &nan_partial, 300);
513 for (nan_partial[0..300]) |value| oracleFold(digest, @as(u32, @bitCast(value)));
514 var nan_selected = nan_values;
515 try select(f32, .ascending, &nan_selected, 519);
516 var nans: u64 = 0;
517 var infinities: u64 = 0;
518 for (nan_selected) |value| {
519 nans += @intFromBool(std.math.isNan(value));
520 infinities += @intFromBool(std.math.isInf(value));
521 }
522 oracleFold(digest, @intFromBool(std.math.isNan(nan_selected[519])));
523 oracleFold(digest, nans);
524 oracleFold(digest, infinities);
525 }
526
527 test "pinned Highway VQSort public APIs agree across full partial select and NaN" {
528 var digest: u64 = 0xcbf2_9ce4_8422_2325;
529 foldIntegerOracle(&digest);
530 try std.testing.expectEqual(@as(u64, 2_625_724_195_073_583_756), digest);
531 foldFloatOracle(&digest);
532 try std.testing.expectEqual(@as(u64, 14_271_760_640_035_470_920), digest);
533 try foldWideOracle(&digest);
534 try std.testing.expectEqual(@as(u64, 8_362_152_917_260_480_414), digest);
535 try foldPartialSelectOracle(&digest);
536 try std.testing.expectEqual(@as(u64, 6_678_231_707_808_607_537), digest);
537 }