lib/simd/src/sort/network.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const constants = @import("constants.zig");
3 const key = @import("key.zig");
4
5 pub const Error = error{
6 TooManyKeys,
7 InvalidKey,
8 };
9
10 pub const Scratch = struct {
11 bytes: [constants.maximum_buffer_bytes]u8 align(64) = undefined,
12 };
13
14 pub fn sort(
15 comptime Key: type,
16 comptime direction: key.Direction,
17 values: []Key,
18 scratch: *Scratch,
19 ) Error!void {
20 const capacity = comptime constants.baseCaseKeyCapacity(Key);
21 if (values.len > capacity) return error.TooManyKeys;
22 for (values) |value| {
23 if (!key.validForNetwork(Key, value)) return error.InvalidKey;
24 }
25 if (values.len < 2) return;
26 const padded = powerOfTwoCeiling(values.len, capacity);
27 const buffer = scratchBuffer(Key, capacity, scratch);
28 @memcpy(buffer[0..values.len], values);
29 @memset(buffer[values.len..padded], key.lastValue(Key, direction));
30 bitonic(Key, direction, buffer[0..padded]);
31 @memcpy(values, buffer[0..values.len]);
32 }
33
34 fn scratchBuffer(
35 comptime Key: type,
36 comptime capacity: usize,
37 scratch: *Scratch,
38 ) *[capacity]Key {
39 comptime std.debug.assert(capacity * @sizeOf(Key) <= constants.maximum_buffer_bytes);
40 return @ptrCast(&scratch.bytes);
41 }
42
43 fn powerOfTwoCeiling(count: usize, capacity: usize) usize {
44 std.debug.assert(count >= 2);
45 std.debug.assert(count <= capacity);
46 var result: usize = 2;
47 var iterations: usize = 0;
48 while (result < count) : (iterations += 1) {
49 std.debug.assert(iterations < 8);
50 result *= 2;
51 }
52 std.debug.assert(result <= capacity);
53 return result;
54 }
55
56 fn bitonic(
57 comptime Key: type,
58 comptime direction: key.Direction,
59 values: []Key,
60 ) void {
61 std.debug.assert(values.len >= 2);
62 std.debug.assert(std.math.isPowerOfTwo(values.len));
63 var span: usize = 2;
64 var span_iterations: usize = 0;
65 while (span <= values.len) : (span *= 2) {
66 std.debug.assert(span_iterations < 8);
67 var distance = span / 2;
68 var distance_iterations: usize = 0;
69 while (distance != 0) : (distance /= 2) {
70 std.debug.assert(distance_iterations < 8);
71 for (0..values.len) |left| {
72 const right = left ^ distance;
73 if (right > left) compareExchange(
74 Key,
75 direction,
76 values,
77 left,
78 right,
79 left & span == 0,
80 );
81 }
82 distance_iterations += 1;
83 }
84 span_iterations += 1;
85 if (span == values.len) break;
86 }
87 }
88
89 fn compareExchange(
90 comptime Key: type,
91 comptime direction: key.Direction,
92 values: []Key,
93 left: usize,
94 right: usize,
95 forward: bool,
96 ) void {
97 std.debug.assert(left < right);
98 std.debug.assert(right < values.len);
99 const a = values[left];
100 const b = values[right];
101 const swap = if (forward)
102 key.networkBefore(Key, direction, b, a)
103 else
104 key.networkBefore(Key, direction, a, b);
105 if (swap) {
106 values[left] = b;
107 values[right] = a;
108 }
109 }
110
111 fn Order(comptime Key: type, comptime direction: key.Direction) type {
112 return struct {
113 fn lessThan(_: void, a: Key, b: Key) bool {
114 return key.networkBefore(Key, direction, a, b);
115 }
116 };
117 }
118
119 fn fillInteger(comptime Key: type, values: []Key) void {
120 const UInt = @Int(.unsigned, @bitSizeOf(Key));
121 var state: u64 = 0x9e37_79b9_7f4a_7c15;
122 for (values, 0..) |*value, index| {
123 state ^= state >> 12;
124 state ^= state << 25;
125 state ^= state >> 27;
126 const mixed = state *% 0x2545_f491_4f6c_dd1d +% index;
127 const bits: UInt = @truncate(mixed);
128 value.* = @bitCast(bits);
129 }
130 }
131
132 fn verifyIntegerType(comptime Key: type) !void {
133 const capacity = comptime constants.baseCaseKeyCapacity(Key);
134 const lengths = [_]usize{ 0, 1, 3, capacity / 2 + 1, capacity - 1, capacity };
135 inline for (.{ key.Direction.ascending, key.Direction.descending }) |direction| {
136 for (lengths) |length| {
137 var storage: [capacity + 2]Key = undefined;
138 fillInteger(Key, storage[1 .. length + 1]);
139 storage[0] = @bitCast(@as(@Int(.unsigned, @bitSizeOf(Key)), 0xa55a));
140 storage[length + 1] = storage[0];
141 var expected: [capacity]Key = undefined;
142 @memcpy(expected[0..length], storage[1 .. length + 1]);
143 std.mem.sort(Key, expected[0..length], {}, Order(Key, direction).lessThan);
144 var scratch = Scratch{};
145 try sort(Key, direction, storage[1 .. length + 1], &scratch);
146 try std.testing.expectEqualSlices(Key, expected[0..length], storage[1 .. length + 1]);
147 try std.testing.expectEqual(storage[0], storage[length + 1]);
148 }
149 }
150 }
151
152 test "Highway VQSort scratch capacity rejects before mutation" {
153 var values = @as([(constants.baseCaseKeyCapacity(u64) + 1)]u64, @splat(17));
154 const expected = values;
155 var scratch = Scratch{};
156 try std.testing.expectError(error.TooManyKeys, sort(u64, .ascending, &values, &scratch));
157 try std.testing.expectEqualSlices(u64, &expected, &values);
158 try std.testing.expectEqual(@as(usize, 1664), @sizeOf(Scratch));
159 try std.testing.expectEqual(@as(usize, 64), @alignOf(Scratch));
160 }
161
162 test "Highway VQSort base networks preserve red zones and order" {
163 try verifyIntegerType(u16);
164 try verifyIntegerType(i16);
165 try verifyIntegerType(u32);
166 try verifyIntegerType(i32);
167 try verifyIntegerType(u64);
168 try verifyIntegerType(i64);
169 }
170
171 test "Highway VQSort base networks sort records and 128-bit keys" {
172 const count = comptime constants.baseCaseKeyCapacity(key.K32V32);
173 var values: [count]key.K32V32 = undefined;
174 for (&values, 0..) |*value, index| {
175 value.* = key.K32V32.init(@intCast((count - index) % 17), @intCast(index));
176 }
177 var expected = values;
178 std.mem.sort(key.K32V32, &expected, {}, Order(key.K32V32, .ascending).lessThan);
179 var scratch = Scratch{};
180 try sort(key.K32V32, .ascending, &values, &scratch);
181 try std.testing.expectEqualSlices(key.K32V32, &expected, &values);
182
183 var wide: [constants.baseCaseKeyCapacity(u128)]u128 = undefined;
184 fillInteger(u128, &wide);
185 var wide_expected = wide;
186 std.mem.sort(u128, &wide_expected, {}, Order(u128, .descending).lessThan);
187 try sort(u128, .descending, &wide, &scratch);
188 try std.testing.expectEqualSlices(u128, &wide_expected, &wide);
189 }
190
191 fn verifyFloatType(comptime Float: type) !void {
192 const capacity = comptime constants.baseCaseKeyCapacity(Float);
193 var values: [capacity]Float = undefined;
194 for (&values, 0..) |*value, index| {
195 const magnitude: Float = @floatFromInt(index + 1);
196 value.* = if (index & 1 == 0) magnitude else -magnitude;
197 }
198 values[capacity / 3] = std.math.inf(Float);
199 values[capacity * 2 / 3] = -std.math.inf(Float);
200 var expected = values;
201 std.mem.sort(Float, &expected, {}, Order(Float, .ascending).lessThan);
202 var scratch = Scratch{};
203 try sort(Float, .ascending, &values, &scratch);
204 try std.testing.expectEqualSlices(Float, &expected, &values);
205 }
206
207 test "Highway VQSort base networks retain floating infinities" {
208 try verifyFloatType(f16);
209 try verifyFloatType(f32);
210 try verifyFloatType(f64);
211 }
212
213 test "Highway VQSort base networks reject NaN before mutation" {
214 var values = [_]f32{ 3, std.math.nan(f32), 1 };
215 const first_bits: u32 = @bitCast(values[0]);
216 const nan_bits: u32 = @bitCast(values[1]);
217 const last_bits: u32 = @bitCast(values[2]);
218 var scratch = Scratch{};
219 try std.testing.expectError(error.InvalidKey, sort(f32, .ascending, &values, &scratch));
220 try std.testing.expectEqual(first_bits, @as(u32, @bitCast(values[0])));
221 try std.testing.expectEqual(nan_bits, @as(u32, @bitCast(values[1])));
222 try std.testing.expectEqual(last_bits, @as(u32, @bitCast(values[2])));
223 }