lib/simd/src/phast/family.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const simd = @import("../root.zig");
3 const hash_mod = simd.hash;
4 const multiply = simd.multiply;
5 const shift = simd.shift;
6
7 pub const PhastError = error{
8 BuildFailed,
9 CapacityExceeded,
10 DuplicateKey,
11 InputOutputOverlap,
12 InvalidData,
13 PlanMismatch,
14 ScratchTooSmall,
15 SeedsTooSmall,
16 };
17
18 pub const Placement = struct {
19 num_slice_offsets: u32 = 0,
20 slice_mask: u32 = 0,
21
22 pub fn init(num_slots: usize, slice_length: usize) PhastError!Placement {
23 if (slice_length == 0 or !std.math.isPowerOfTwo(slice_length)) {
24 return error.InvalidData;
25 }
26 if (num_slots < slice_length or num_slots > std.math.maxInt(u32)) {
27 return error.CapacityExceeded;
28 }
29 if (slice_length > std.math.maxInt(u32)) return error.CapacityExceeded;
30 const num_slice_offsets = num_slots - slice_length + 1;
31 return .{
32 .num_slice_offsets = @intCast(num_slice_offsets),
33 .slice_mask = @intCast(slice_length - 1),
34 };
35 }
36
37 pub fn sliceLength(self: Placement) usize {
38 return @as(usize, self.slice_mask) + 1;
39 }
40 };
41
42 pub const Config = struct {
43 num_slots: usize = 0,
44 hash_key: u32 = 0,
45 bucket_mask: u32 = 0,
46 placement: Placement = .{},
47
48 pub fn numBuckets(self: Config) usize {
49 return @as(usize, self.bucket_mask) + 1;
50 }
51
52 pub fn allocatedBytes(self: Config, payload_bytes: usize) PhastError!usize {
53 const payload = std.math.mul(usize, self.num_slots, payload_bytes) catch
54 return error.CapacityExceeded;
55 const seed_bytes = std.mem.alignForward(usize, self.numBuckets(), @sizeOf(u32));
56 return std.math.add(usize, payload, seed_bytes) catch
57 return error.CapacityExceeded;
58 }
59
60 pub fn seedWordsLen(self: Config) usize {
61 return seedWordLength(self.numBuckets());
62 }
63 };
64
65 pub const Seeds = struct {
66 words: []const u32 = &.{},
67
68 pub fn get(self: Seeds, bucket_index: u32) u32 {
69 std.debug.assert(@as(usize, bucket_index) < self.words.len * @sizeOf(u32));
70 const bit_index = (bucket_index & 3) * 8;
71 return (self.words[bucket_index >> 2] >> @intCast(bit_index)) & 0xff;
72 }
73 };
74
75 pub const Data = struct {
76 config: Config = .{},
77 seeds: Seeds = .{},
78 config_index: usize = 0,
79 attempt_index: usize = 0,
80
81 pub fn numSlots(self: Data) usize {
82 return self.config.num_slots;
83 }
84
85 pub fn allocatedBytes(self: Data, payload_bytes: usize) PhastError!usize {
86 return self.config.allocatedBytes(payload_bytes);
87 }
88
89 pub fn isEmpty(self: Data) bool {
90 return self.config.num_slots == 0;
91 }
92 };
93
94 pub const Phast = struct {
95 data: Data,
96 hash: hash_mod.Triple32,
97
98 const Self = @This();
99
100 pub fn init(data: Data) PhastError!Self {
101 if (data.isEmpty()) return error.BuildFailed;
102 const num_buckets = data.config.numBuckets();
103 const slice_length = data.config.placement.sliceLength();
104 if (data.config.num_slots > std.math.maxInt(u32) or
105 num_buckets > std.math.maxInt(u32) or
106 !std.math.isPowerOfTwo(num_buckets) or
107 data.seeds.words.len != data.config.seedWordsLen() or
108 data.config.placement.num_slice_offsets == 0 or
109 !std.math.isPowerOfTwo(slice_length) or
110 slice_length > data.config.num_slots or
111 data.config.placement.num_slice_offsets !=
112 data.config.num_slots - slice_length + 1)
113 {
114 return error.InvalidData;
115 }
116 return .{
117 .data = data,
118 .hash = hash_mod.Triple32.initKey(data.config.hash_key),
119 };
120 }
121
122 pub fn index(self: Self, key: u32) u32 {
123 const hash = self.hash.hash(key);
124 const seed = self.data.seeds.get(hash & self.data.config.bucket_mask);
125 return positionFromHashAndSeed(self.data.config.placement, hash, seed);
126 }
127
128 pub fn twoVec(
129 self: Self,
130 comptime D: type,
131 first_keys: D.Vector,
132 second_keys: D.Vector,
133 ) VectorPair(D) {
134 requireTag(D);
135 var first_hashes = first_keys;
136 var second_hashes = second_keys;
137 self.hash.twoVec(D, &first_hashes, &second_hashes);
138 var first_seeds: D.Vector = undefined;
139 var second_seeds: D.Vector = undefined;
140 inline for (0..D.lane_count) |lane| {
141 first_seeds[lane] = self.data.seeds.get(
142 first_hashes[lane] & self.data.config.bucket_mask,
143 );
144 second_seeds[lane] = self.data.seeds.get(
145 second_hashes[lane] & self.data.config.bucket_mask,
146 );
147 }
148 return positionPairFromHashesAndSeeds(
149 self.data.config.placement,
150 D,
151 first_hashes,
152 second_hashes,
153 first_seeds,
154 second_seeds,
155 );
156 }
157
158 pub fn indexBatch(
159 self: Self,
160 comptime D: type,
161 keys: []const u32,
162 indices: []u32,
163 ) PhastError!void {
164 requireTag(D);
165 if (indices.len != keys.len) return error.PlanMismatch;
166 var offset: usize = 0;
167 while (offset + 2 * D.lane_count <= keys.len) : (offset += 2 * D.lane_count) {
168 const first: D.Vector = keys[offset..][0..D.lane_count].*;
169 const second: D.Vector = keys[offset + D.lane_count ..][0..D.lane_count].*;
170 const result = self.twoVec(D, first, second);
171 const first_indices: [D.lane_count]u32 = result.first;
172 const second_indices: [D.lane_count]u32 = result.second;
173 @memcpy(indices[offset..][0..D.lane_count], &first_indices);
174 @memcpy(indices[offset + D.lane_count ..][0..D.lane_count], &second_indices);
175 }
176 for (keys[offset..], indices[offset..]) |key, *output| output.* = self.index(key);
177 }
178 };
179
180 pub fn VectorPair(comptime D: type) type {
181 requireTag(D);
182 return struct {
183 first: D.Vector,
184 second: D.Vector,
185 };
186 }
187
188 pub fn positionFromHashAndSeed(placement: Placement, hash: u32, seed: u32) u32 {
189 std.debug.assert(seed < 256);
190 const slice_offset = hash_mod.lemireMod(hash, placement.num_slice_offsets);
191 const within_slice = @as(u32, hash16(@truncate((hash >> 16) +% seed))) &
192 placement.slice_mask;
193 return slice_offset + within_slice;
194 }
195
196 pub fn seedWordLength(num_buckets: usize) usize {
197 return std.math.divCeil(usize, num_buckets, @sizeOf(u32)) catch unreachable;
198 }
199
200 pub fn positionPairFromHashesAndSeeds(
201 placement: Placement,
202 comptime D: type,
203 first_hashes: D.Vector,
204 second_hashes: D.Vector,
205 first_seeds: D.Vector,
206 second_seeds: D.Vector,
207 ) VectorPair(D) {
208 const D16 = D.repartition(u16);
209 var combined: D16.Vector = undefined;
210 inline for (0..D.lane_count) |lane| {
211 combined[lane] = @as(u16, @truncate(first_hashes[lane] >> 16)) +%
212 @as(u16, @truncate(first_seeds[lane]));
213 combined[D.lane_count + lane] =
214 @as(u16, @truncate(second_hashes[lane] >> 16)) +%
215 @as(u16, @truncate(second_seeds[lane]));
216 }
217 const within = hash16Vector(D16, combined);
218 const offsets: D.Vector = @splat(placement.num_slice_offsets);
219 const mask: D.Vector = @splat(placement.slice_mask);
220 var result: VectorPair(D) = .{
221 .first = multiply.mulHigh(D, first_hashes, offsets),
222 .second = multiply.mulHigh(D, second_hashes, offsets),
223 };
224 inline for (0..D.lane_count) |lane| {
225 result.first[lane] += @as(u32, within[lane]) & mask[lane];
226 result.second[lane] += @as(u32, within[D.lane_count + lane]) & mask[lane];
227 }
228 return result;
229 }
230
231 fn hash16(initial: u16) u16 {
232 var value = initial;
233 value ^= value >> 8;
234 const squared = value *% value;
235 value +%= squared *% 0xca32;
236 value ^= value >> 12;
237 value *%= 0x3929;
238 return value;
239 }
240
241 fn hash16Vector(comptime D: type, initial: D.Vector) D.Vector {
242 var value = initial;
243 value ^= shift.shiftRight(D, 8, value);
244 value +%= (value *% value) *% @as(D.Vector, @splat(0xca32));
245 value ^= shift.shiftRight(D, 12, value);
246 value *%= @splat(0x3929);
247 return value;
248 }
249
250 fn requireTag(comptime D: type) void {
251 if (comptime D.Lane != u32) @compileError("PHAST tags require u32 lanes");
252 }