lib/simd/src/cuckoo2x2.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const hash_mod = @import("hash.zig");
3 const indexed = @import("indexed.zig");
4 const random = @import("random.zig");
5 const shift = @import("shift.zig");
6 const tag = @import("tag.zig");
7
8 pub const minimum_bucket_count: usize = 262_144;
9 pub const configuration_count: usize = 3;
10 pub const max_attempts: usize = 512;
11 pub const max_displacements: usize = 500;
12
13 pub const Cuckoo2x2Config = struct {
14 bucket_mask: u32 = 0,
15 hash_key: u32 = 0,
16
17 const Self = @This();
18
19 pub fn init(num_buckets: usize, hash_key: u32) Self {
20 std.debug.assert(num_buckets != 0);
21 std.debug.assert(num_buckets <= std.math.maxInt(u32));
22 std.debug.assert(std.math.isPowerOfTwo(num_buckets));
23 return .{
24 .bucket_mask = @intCast(num_buckets - 1),
25 .hash_key = hash_key,
26 };
27 }
28
29 pub fn numBuckets(self: Self) usize {
30 return @as(usize, self.bucket_mask) + 1;
31 }
32 };
33
34 pub const Cuckoo2x2Data = struct {
35 config: Cuckoo2x2Config = .{},
36 entries: []const u32 = &.{},
37 config_idx: usize = 0,
38 attempt_idx: usize = 0,
39 num_primary: u32 = 0,
40 num_secondary: u32 = 0,
41
42 const Self = @This();
43
44 pub fn isEmpty(self: Self) bool {
45 return self.entries.len == 0;
46 }
47
48 pub fn numBuckets(self: Self) usize {
49 if (self.isEmpty()) return 0;
50 std.debug.assert(self.entries.len == self.config.numBuckets());
51 return self.entries.len;
52 }
53
54 pub fn allocatedBytes(self: Self) usize {
55 return self.entries.len * @sizeOf(u32);
56 }
57 };
58
59 pub const Cuckoo2x2 = struct {
60 hash1: hash_mod.WeakTwoMul,
61 data: Cuckoo2x2Data,
62
63 const Self = @This();
64
65 pub fn init(data: Cuckoo2x2Data) Self {
66 std.debug.assert(!data.isEmpty());
67 std.debug.assert(data.numBuckets() >= minimum_bucket_count);
68 return .{
69 .hash1 = hash_mod.WeakTwoMul.initKey(data.config.hash_key),
70 .data = data,
71 };
72 }
73
74 pub fn getData(self: Self) Cuckoo2x2Data {
75 return self.data;
76 }
77
78 pub fn contains(self: Self, key: u32) bool {
79 const h1 = self.hash1.hash(key);
80 const b1 = h1 & self.data.config.bucket_mask;
81 const fingerprint = h1 >> 18;
82 const b2 = b1 ^ (fingerprint + 1);
83 const primary: u16 = @truncate(fingerprint | 0x4000);
84 const secondary: u16 = @truncate(fingerprint | 0x8000);
85 const e1 = self.data.entries[b1];
86 const e2 = self.data.entries[b2];
87 return primary == @as(u16, @truncate(e1)) or
88 primary == @as(u16, @truncate(e1 >> 16)) or
89 secondary == @as(u16, @truncate(e2)) or
90 secondary == @as(u16, @truncate(e2 >> 16));
91 }
92
93 pub fn query(self: Self, comptime D: type, keys: D.Vector) D.Mask {
94 requireTag(D);
95 const h1 = self.hash1.oneVec(D, keys);
96 const b1 = h1 & @as(D.Vector, @splat(self.data.config.bucket_mask));
97 const fingerprint = shift.shiftRight(D, 18, h1);
98 const b2 = b1 ^ (fingerprint +% @as(D.Vector, @splat(1)));
99 const primary = fingerprint | @as(D.Vector, @splat(0x4000));
100 const secondary = fingerprint | @as(D.Vector, @splat(0x8000));
101 const e1 = indexed.gatherIndex(D, self.data.entries, b1);
102 const e2 = indexed.gatherIndex(D, self.data.entries, b2);
103 const low_mask: D.Vector = @splat(0xffff);
104 const found = (e1 & low_mask == primary) |
105 (shift.shiftRight(D, 16, e1) == primary) |
106 (e2 & low_mask == secondary) |
107 (shift.shiftRight(D, 16, e2) == secondary);
108 return !found;
109 }
110
111 pub fn containsVec(self: Self, comptime D: type, keys: D.Vector) D.Mask {
112 return !self.query(D, keys);
113 }
114 };
115
116 pub const Cuckoo2x2BuildError = error{
117 BuildFailed,
118 CapacityExceeded,
119 DuplicateKey,
120 EntriesTooSmall,
121 PlanMismatch,
122 ScratchTooSmall,
123 };
124
125 pub const Cuckoo2x2Plan = struct {
126 key_count: usize,
127 bucket_counts: [configuration_count]usize,
128 max_bucket_count: usize,
129 entries_len: usize,
130 scratch_len: usize,
131
132 const Self = @This();
133
134 pub fn inspect(key_count: usize) Cuckoo2x2BuildError!Self {
135 if (key_count > std.math.maxInt(u32)) return error.CapacityExceeded;
136 const quarter = key_count / 4;
137 const base = if (quarter <= 1)
138 @as(usize, 1)
139 else
140 std.math.ceilPowerOfTwo(usize, quarter) catch return error.CapacityExceeded;
141 var bucket_counts: [configuration_count]usize = undefined;
142 inline for (0..configuration_count) |index| {
143 const multiplier = @as(usize, 1) << index;
144 const candidate = std.math.mul(usize, base, multiplier) catch
145 return error.CapacityExceeded;
146 const num_buckets = @max(candidate, minimum_bucket_count);
147 if (num_buckets > std.math.maxInt(u32)) return error.CapacityExceeded;
148 bucket_counts[index] = num_buckets;
149 }
150 const max_bucket_count = bucket_counts[configuration_count - 1];
151 const bucket_key_count = std.math.mul(usize, max_bucket_count, 2) catch
152 return error.CapacityExceeded;
153 const byte_count = std.math.add(usize, key_count, max_bucket_count) catch
154 return error.CapacityExceeded;
155 const byte_words = std.math.divCeil(usize, byte_count, @sizeOf(u32)) catch
156 unreachable;
157 const hash_and_keys = std.math.add(usize, key_count, bucket_key_count) catch
158 return error.CapacityExceeded;
159 const scratch_len = std.math.add(usize, hash_and_keys, byte_words) catch
160 return error.CapacityExceeded;
161 return .{
162 .key_count = key_count,
163 .bucket_counts = bucket_counts,
164 .max_bucket_count = max_bucket_count,
165 .entries_len = max_bucket_count,
166 .scratch_len = scratch_len,
167 };
168 }
169
170 pub fn build(
171 self: Self,
172 scratch: []u32,
173 entries: []u32,
174 keys: []const u32,
175 ) Cuckoo2x2BuildError!Cuckoo2x2Data {
176 if (keys.len != self.key_count) return error.PlanMismatch;
177 if (scratch.len < self.scratch_len) return error.ScratchTooSmall;
178 if (entries.len < self.entries_len) return error.EntriesTooSmall;
179 const parts = self.sections(scratch);
180 const engine = random.AesCtrEngine.initDeterministic();
181 const distinct_hash = hash_mod.WeakTwoMul.initSeed(&engine, 0);
182 hash_mod.hashArray(tag.ScalableTag(u32), distinct_hash, keys, parts.hashes);
183 std.mem.sort(u32, parts.hashes, {}, std.sort.asc(u32));
184 const sorted_hashes = parts.hashes;
185 if (sorted_hashes.len > 1) {
186 const previous_hashes = sorted_hashes[0 .. sorted_hashes.len - 1];
187 for (sorted_hashes[1..], previous_hashes) |current, previous| {
188 if (current == previous) return error.DuplicateKey;
189 }
190 }
191
192 for (self.bucket_counts, 0..) |num_buckets, config_idx| {
193 for (0..max_attempts) |attempt_idx| {
194 const hash_key: u32 = @truncate(engine.generate(attempt_idx, 0));
195 const config = Cuckoo2x2Config.init(num_buckets, hash_key);
196 const hash1 = hash_mod.WeakTwoMul.initKey(hash_key);
197 @memset(parts.counts[0..num_buckets], 0);
198 @memset(entries[0..num_buckets], 0);
199 hash_mod.hashArray(tag.ScalableTag(u32), hash1, keys, parts.hashes);
200 if (!cuckooAssign(
201 config,
202 parts.hashes,
203 parts.choices,
204 parts.counts[0..num_buckets],
205 parts.bucket_keys[0 .. num_buckets * 2],
206 )) continue;
207 populateEntries(
208 config,
209 parts.hashes,
210 parts.choices,
211 parts.counts[0..num_buckets],
212 entries[0..num_buckets],
213 );
214 var num_primary: u32 = 0;
215 for (parts.choices) |choice| num_primary += @intFromBool(choice == 0);
216 return .{
217 .config = config,
218 .entries = entries[0..num_buckets],
219 .config_idx = config_idx,
220 .attempt_idx = attempt_idx,
221 .num_primary = num_primary,
222 .num_secondary = @as(u32, @intCast(keys.len)) - num_primary,
223 };
224 }
225 }
226 return error.BuildFailed;
227 }
228
229 fn sections(self: Self, scratch: []u32) ScratchSections {
230 const hashes = scratch[0..self.key_count];
231 const bucket_keys_begin = self.key_count;
232 const bucket_keys_end = bucket_keys_begin + self.max_bucket_count * 2;
233 const bucket_keys = scratch[bucket_keys_begin..bucket_keys_end];
234 const bytes = std.mem.sliceAsBytes(scratch[bucket_keys_end..self.scratch_len]);
235 const choices = bytes[0..self.key_count];
236 const counts = bytes[self.key_count .. self.key_count + self.max_bucket_count];
237 return .{
238 .hashes = hashes,
239 .bucket_keys = bucket_keys,
240 .choices = choices,
241 .counts = counts,
242 };
243 }
244 };
245
246 pub fn cuckoo2x2ScratchLen(key_count: usize) Cuckoo2x2BuildError!usize {
247 return (try Cuckoo2x2Plan.inspect(key_count)).scratch_len;
248 }
249
250 pub fn cuckoo2x2EntriesLen(key_count: usize) Cuckoo2x2BuildError!usize {
251 return (try Cuckoo2x2Plan.inspect(key_count)).entries_len;
252 }
253
254 pub fn buildCuckoo2x2(
255 scratch: []u32,
256 entries: []u32,
257 keys: []const u32,
258 ) Cuckoo2x2BuildError!Cuckoo2x2Data {
259 const plan = try Cuckoo2x2Plan.inspect(keys.len);
260 return plan.build(scratch, entries, keys);
261 }
262
263 pub fn makeCuckoo2x2(
264 scratch: []u32,
265 entries: []u32,
266 keys: []const u32,
267 ) Cuckoo2x2BuildError!Cuckoo2x2 {
268 return Cuckoo2x2.init(try buildCuckoo2x2(scratch, entries, keys));
269 }
270
271 const ScratchSections = struct {
272 hashes: []u32,
273 bucket_keys: []u32,
274 choices: []u8,
275 counts: []u8,
276 };
277
278 fn bucket(config: Cuckoo2x2Config, hashes: []const u32, key_idx: usize, which: u8) usize {
279 const value = hashes[key_idx];
280 const primary = value & config.bucket_mask;
281 if (which == 0) return primary;
282 return primary ^ ((value >> 18) + 1);
283 }
284
285 fn cuckooAssign(
286 config: Cuckoo2x2Config,
287 hashes: []const u32,
288 choices: []u8,
289 counts: []u8,
290 bucket_keys: []u32,
291 ) bool {
292 std.debug.assert(hashes.len == choices.len);
293 std.debug.assert(counts.len == config.numBuckets());
294 std.debug.assert(bucket_keys.len == counts.len * 2);
295 @memset(bucket_keys, std.math.maxInt(u32));
296 for (0..hashes.len) |key_idx| {
297 var displaced = key_idx;
298 const primary = bucket(config, hashes, displaced, 0);
299 const secondary = bucket(config, hashes, displaced, 1);
300 if (counts[primary] <= counts[secondary] and counts[primary] < 2) {
301 choices[displaced] = 0;
302 bucket_keys[primary * 2 + counts[primary]] = @intCast(displaced);
303 counts[primary] += 1;
304 continue;
305 }
306 if (counts[secondary] < 2) {
307 choices[displaced] = 1;
308 bucket_keys[secondary * 2 + counts[secondary]] = @intCast(displaced);
309 counts[secondary] += 1;
310 continue;
311 }
312 if (counts[primary] < 2) {
313 choices[displaced] = 0;
314 bucket_keys[primary * 2 + counts[primary]] = @intCast(displaced);
315 counts[primary] += 1;
316 continue;
317 }
318
319 var which: u8 = 0;
320 var placed = false;
321 for (0..max_displacements) |_| {
322 const destination = bucket(config, hashes, displaced, which);
323 if (counts[destination] < 2) {
324 choices[displaced] = which;
325 bucket_keys[destination * 2 + counts[destination]] = @intCast(displaced);
326 counts[destination] += 1;
327 placed = true;
328 break;
329 }
330 const evicted = bucket_keys[destination * 2];
331 bucket_keys[destination * 2] = bucket_keys[destination * 2 + 1];
332 bucket_keys[destination * 2 + 1] = @intCast(displaced);
333 choices[displaced] = which;
334 displaced = evicted;
335 which = 1 - choices[evicted];
336 }
337 if (!placed) return false;
338 }
339 return true;
340 }
341
342 fn populateEntries(
343 config: Cuckoo2x2Config,
344 hashes: []const u32,
345 choices: []const u8,
346 counts: []u8,
347 entries: []u32,
348 ) void {
349 std.debug.assert(hashes.len == choices.len);
350 std.debug.assert(counts.len == entries.len);
351 @memset(counts, 0);
352 for (hashes, choices) |value, choice| {
353 const primary = value & config.bucket_mask;
354 const fingerprint = value >> 18;
355 const destination = if (choice == 0) primary else primary ^ (fingerprint + 1);
356 const tagged = fingerprint | if (choice == 0) @as(u32, 0x4000) else 0x8000;
357 const slot = counts[destination];
358 std.debug.assert(slot < 2);
359 entries[destination] |= tagged << @intCast(slot * 16);
360 counts[destination] += 1;
361 }
362 for (entries, counts) |*entry, count| {
363 if (count == 1) entry.* |= entry.* << 16;
364 }
365 }
366
367 fn requireTag(comptime D: type) void {
368 if (comptime D.Lane != u32) @compileError("Cuckoo2x2 requires u32 lanes");
369 }
370
371 test "Highway Cuckoo2x2 scalar and vector membership agree" {
372 const key_count = 1000;
373 var keys: [key_count]u32 = undefined;
374 for (&keys, 0..) |*key, index| key.* = @as(u32, @intCast(index)) *% 37 +% 1;
375 const plan = try Cuckoo2x2Plan.inspect(keys.len);
376 const allocator = std.testing.allocator;
377 const scratch = try allocator.alloc(u32, plan.scratch_len);
378 defer allocator.free(scratch);
379 const entries = try allocator.alloc(u32, plan.entries_len);
380 defer allocator.free(entries);
381 const set = Cuckoo2x2.init(try plan.build(scratch, entries, &keys));
382 for (keys) |key| try std.testing.expect(set.contains(key));
383 const D = tag.FixedTag(u32, 8);
384 var index: usize = 0;
385 while (index + D.lane_count <= keys.len) : (index += D.lane_count) {
386 const input: D.Vector = keys[index..][0..D.lane_count].*;
387 try std.testing.expect(@reduce(.And, set.containsVec(D, input)));
388 }
389 for (0..key_count) |absent_idx| {
390 const absent = @as(u32, @intCast(absent_idx)) *% 37 +% 2;
391 try std.testing.expect(!set.contains(absent));
392 }
393 }
394
395 test "Highway Cuckoo2x2 single-worker builder oracle matches exactly" {
396 const key_count = 1024;
397 var keys: [key_count]u32 = undefined;
398 for (&keys, 0..) |*key, index| key.* = @as(u32, @intCast(index)) *% 37 +% 1;
399 const plan = try Cuckoo2x2Plan.inspect(keys.len);
400 const allocator = std.testing.allocator;
401 const scratch = try allocator.alloc(u32, plan.scratch_len);
402 defer allocator.free(scratch);
403 const entries = try allocator.alloc(u32, plan.entries_len);
404 defer allocator.free(entries);
405 const data = try plan.build(scratch, entries, &keys);
406 try std.testing.expectEqual(@as(usize, 262_144), data.numBuckets());
407 try std.testing.expectEqual(@as(u32, 0x7aa5_332d), data.config.hash_key);
408 try std.testing.expectEqual(@as(usize, 0), data.config_idx);
409 try std.testing.expectEqual(@as(usize, 0), data.attempt_idx);
410 try std.testing.expectEqual(@as(u32, 1023), data.num_primary);
411 try std.testing.expectEqual(@as(u32, 1), data.num_secondary);
412 const expected_indices = [16]usize{
413 90, 153, 321, 423, 429, 636, 747, 1234,
414 1648, 2094, 2308, 2564, 2586, 2598, 3034, 4066,
415 };
416 const expected_entries = [16]u32{
417 0x6dd4_6dd4,
418 0x5894_5894,
419 0x686c_686c,
420 0x5272_5272,
421 0x454f_454f,
422 0x6a0e_6a0e,
423 0x4a8e_4a8e,
424 0x745b_745b,
425 0x4499_4499,
426 0x503c_503c,
427 0x4800_4800,
428 0x57f3_57f3,
429 0x6dd1_6dd1,
430 0x7c93_7c93,
431 0x544e_544e,
432 0x5634_5634,
433 };
434 for (expected_indices, expected_entries) |index, expected| {
435 try std.testing.expectEqual(expected, data.entries[index]);
436 }
437 var digest: u64 = 1_469_598_103_934_665_603;
438 var nonzero: usize = 0;
439 for (data.entries) |entry| {
440 digest = (digest ^ entry) *% 1_099_511_628_211;
441 nonzero += @intFromBool(entry != 0);
442 }
443 try std.testing.expectEqual(@as(u64, 0x947b_8ca0_29af_1476), digest);
444 try std.testing.expectEqual(key_count, nonzero);
445 }
446
447 test "Highway Cuckoo2x2 builds and queries upstream sizes" {
448 const allocator = std.testing.allocator;
449 for ([_]usize{ 6000, 60_000 }) |key_count| {
450 const keys = try allocator.alloc(u32, key_count);
451 defer allocator.free(keys);
452 hash_mod.fillRandomDistinct(u32, 0, keys);
453 const plan = try Cuckoo2x2Plan.inspect(keys.len);
454 const scratch = try allocator.alloc(u32, plan.scratch_len);
455 defer allocator.free(scratch);
456 const entries = try allocator.alloc(u32, plan.entries_len);
457 defer allocator.free(entries);
458 const data = try plan.build(scratch, entries, keys);
459 const set = Cuckoo2x2.init(data);
460 try std.testing.expectEqual(key_count, data.num_primary + data.num_secondary);
461 try std.testing.expect(data.numBuckets() >= minimum_bucket_count);
462 for (keys) |key| try std.testing.expect(set.contains(key));
463 }
464 }
465
466 test "Highway Cuckoo2x2 plans enforce capacities and distinct keys" {
467 const keys = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 };
468 const plan = try Cuckoo2x2Plan.inspect(keys.len);
469 const allocator = std.testing.allocator;
470 const scratch = try allocator.alloc(u32, plan.scratch_len);
471 defer allocator.free(scratch);
472 const entries = try allocator.alloc(u32, plan.entries_len);
473 defer allocator.free(entries);
474 try std.testing.expectError(
475 error.ScratchTooSmall,
476 plan.build(scratch[0 .. scratch.len - 1], entries, &keys),
477 );
478 try std.testing.expectError(
479 error.EntriesTooSmall,
480 plan.build(scratch, entries[0 .. entries.len - 1], &keys),
481 );
482 try std.testing.expectError(error.PlanMismatch, plan.build(scratch, entries, keys[0..7]));
483 const duplicate = [_]u32{ 1, 2, 3, 3, 5, 6, 7, 8 };
484 try std.testing.expectError(error.DuplicateKey, plan.build(scratch, entries, &duplicate));
485 }
486
487 test "Highway Cuckoo2x2 represents the empty set" {
488 const plan = try Cuckoo2x2Plan.inspect(0);
489 const allocator = std.testing.allocator;
490 const scratch = try allocator.alloc(u32, plan.scratch_len);
491 defer allocator.free(scratch);
492 const entries = try allocator.alloc(u32, plan.entries_len);
493 defer allocator.free(entries);
494 const data = try plan.build(scratch, entries, &.{});
495 const set = Cuckoo2x2.init(data);
496 try std.testing.expect(!set.contains(0));
497 try std.testing.expect(!set.contains(std.math.maxInt(u32)));
498 try std.testing.expectEqual(@as(u32, 0), data.num_primary);
499 try std.testing.expectEqual(@as(u32, 0), data.num_secondary);
500 }