lib/simd/src/table.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn tableLookupBytes(comptime D: type, table: anytype, indices: D.Vector) D.Vector {
4 return lookupBytes(D, false, table, indices);
5 }
6
7 pub fn tableLookupBytesOr0(comptime D: type, table: anytype, indices: D.Vector) D.Vector {
8 return lookupBytes(D, true, table, indices);
9 }
10
11 pub fn bitShuffle(comptime D: type, values: D.Vector, indices: anytype) D.Vector {
12 if (D.Lane != u64 and D.Lane != i64) @compileError("bitShuffle requires 64-bit integer lanes");
13 const index_info = vectorInfo(@TypeOf(indices));
14 if (@typeInfo(index_info.child) != .int or @bitSizeOf(index_info.child) != 8) {
15 @compileError("bitShuffle indices require 8-bit integer lanes");
16 }
17 if (index_info.len != D.lane_count * 8) @compileError("bitShuffle requires eight indices per lane");
18 const value_lanes: [D.lane_count]D.Lane = @bitCast(values);
19 const index_lanes: [index_info.len]index_info.child = @bitCast(indices);
20 var result: [D.lane_count]u64 = @splat(0);
21 inline for (0..D.lane_count) |lane| {
22 const bits: u64 = @bitCast(value_lanes[lane]);
23 inline for (0..8) |bit| {
24 const index = checkedIndex(index_lanes[lane * 8 + bit]);
25 std.debug.assert(index < 64);
26 result[lane] |= ((bits >> @intCast(index)) & 1) << bit;
27 }
28 }
29 return @bitCast(result);
30 }
31
32 pub fn indicesFromVec(comptime D: type, indices: anytype) @TypeOf(indices) {
33 const info = vectorInfo(@TypeOf(indices));
34 if (info.len < D.lane_count or @typeInfo(info.child) != .int or
35 @bitSizeOf(info.child) != @bitSizeOf(D.Lane))
36 {
37 @compileError("table indices require enough integer lanes matching the table lane width");
38 }
39 return indices;
40 }
41
42 pub fn setTableIndices(
43 comptime D: type,
44 input: []const indexLane(D),
45 ) D.rebind(indexLane(D)).Vector {
46 std.debug.assert(input.len >= D.lane_count);
47 return @as(D.rebind(indexLane(D)).Vector, input[0..D.lane_count].*);
48 }
49
50 pub fn tableLookupLanes(comptime D: type, table: D.Vector, indices: anytype) D.Vector {
51 const index_info = validateLaneIndices(D, @TypeOf(indices));
52 const table_lanes: [D.lane_count]D.Lane = @bitCast(table);
53 const index_lanes: [index_info.len]index_info.child = @bitCast(indices);
54 var result: [D.lane_count]D.Lane = undefined;
55 inline for (0..D.lane_count) |lane| {
56 const index = checkedIndex(index_lanes[lane]);
57 std.debug.assert(index < D.lane_count);
58 result[lane] = table_lanes[index];
59 }
60 return @bitCast(result);
61 }
62
63 pub fn twoTablesLookupLanes(
64 comptime D: type,
65 first: D.Vector,
66 second: D.Vector,
67 indices: anytype,
68 ) D.Vector {
69 const index_info = validateLaneIndices(D, @TypeOf(indices));
70 const first_lanes: [D.lane_count]D.Lane = @bitCast(first);
71 const second_lanes: [D.lane_count]D.Lane = @bitCast(second);
72 const index_lanes: [index_info.len]index_info.child = @bitCast(indices);
73 var result: [D.lane_count]D.Lane = undefined;
74 inline for (0..D.lane_count) |lane| {
75 const index = checkedIndex(index_lanes[lane]);
76 std.debug.assert(index < D.lane_count * 2);
77 result[lane] = if (index < D.lane_count)
78 first_lanes[index]
79 else
80 second_lanes[index - D.lane_count];
81 }
82 return @bitCast(result);
83 }
84
85 pub fn lookup8(comptime D: type, table: []const D.Lane, indices: anytype) D.Vector {
86 return lookupN(D, 8, table, indices);
87 }
88
89 pub fn lookup16(comptime D: type, table: []const D.Lane, indices: anytype) D.Vector {
90 return lookupN(D, 16, table, indices);
91 }
92
93 pub fn lookup32(comptime D: type, table: []const D.Lane, indices: anytype) D.Vector {
94 return lookupN(D, 32, table, indices);
95 }
96
97 pub fn canLookup8(comptime D: type) bool {
98 _ = D;
99 return true;
100 }
101
102 pub fn canLookup16(comptime D: type) bool {
103 _ = D;
104 return true;
105 }
106
107 pub fn canLookup32(comptime D: type) bool {
108 _ = D;
109 return true;
110 }
111
112 fn lookupBytes(comptime D: type, comptime or_zero: bool, table: anytype, indices: D.Vector) D.Vector {
113 if (@typeInfo(D.Lane) != .int) @compileError("byte table indices require integer vectors");
114 const table_type = @TypeOf(table);
115 const table_info = vectorInfo(table_type);
116 if (@typeInfo(table_info.child) != .int) @compileError("byte lookup tables require integer vectors");
117 const table_byte_count = @bitSizeOf(table_type) / 8;
118 if (table_byte_count == 0) @compileError("byte lookup table cannot be empty");
119 const table_bytes: [table_byte_count]u8 = @bitCast(table);
120 const index_bytes: [D.byte_count]u8 = @bitCast(indices);
121 var result: [D.byte_count]u8 = undefined;
122 inline for (0..D.byte_count) |index| {
123 const requested = index_bytes[index];
124 if (or_zero and requested & 0x80 != 0) {
125 result[index] = 0;
126 } else {
127 const block = index / 16 * 16;
128 result[index] = table_bytes[(block + requested) % @min(table_byte_count, 256)];
129 }
130 }
131 return @bitCast(result);
132 }
133
134 fn lookupN(
135 comptime D: type,
136 comptime table_size: usize,
137 table: []const D.Lane,
138 indices: anytype,
139 ) D.Vector {
140 const index_info = validateLaneIndices(D, @TypeOf(indices));
141 std.debug.assert(table.len >= table_size);
142 const index_lanes: [index_info.len]index_info.child = @bitCast(indices);
143 var result: [D.lane_count]D.Lane = undefined;
144 inline for (0..D.lane_count) |lane| {
145 const index = checkedIndex(index_lanes[lane]);
146 std.debug.assert(index < table_size);
147 result[lane] = table[index];
148 }
149 return @bitCast(result);
150 }
151
152 fn validateLaneIndices(comptime D: type, comptime T: type) @TypeOf(vectorInfo(T)) {
153 const info = vectorInfo(T);
154 if (@typeInfo(info.child) != .int or info.len < D.lane_count or
155 @bitSizeOf(info.child) != @bitSizeOf(D.Lane))
156 {
157 @compileError("lane lookup indices require enough same-width integer lanes");
158 }
159 return info;
160 }
161
162 fn vectorInfo(comptime T: type) std.builtin.Type.Vector {
163 return switch (@typeInfo(T)) {
164 .vector => |info| info,
165 else => @compileError("table operation requires a vector"),
166 };
167 }
168
169 fn checkedIndex(value: anytype) usize {
170 const T = @TypeOf(value);
171 if (comptime @typeInfo(T).int.signedness == .signed) std.debug.assert(value >= 0);
172 return @intCast(value);
173 }
174
175 fn indexLane(comptime D: type) type {
176 return @Int(.unsigned, @bitSizeOf(D.Lane));
177 }
178
179 test "Highway byte table lookup is block-local and optionally zeroes high-bit indices" {
180 const simd = @import("root.zig");
181 const D = simd.FixedTag(u8, 32);
182 const table = simd.iota(D, 1);
183 const indices: D.Vector = .{
184 15, 0, 2, 2, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 1,
185 15, 0, 2, 2, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 1,
186 };
187 try std.testing.expect(@reduce(.And, tableLookupBytes(D, table, indices) == @as(D.Vector, .{
188 16, 1, 3, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 2,
189 32, 17, 19, 19, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 18,
190 })));
191 var zero_indices = indices;
192 zero_indices[3] = 0x82;
193 zero_indices[20] = 0x90;
194 const zeroed = tableLookupBytesOr0(D, table, zero_indices);
195 try std.testing.expectEqual(@as(u8, 0), zeroed[3]);
196 try std.testing.expectEqual(@as(u8, 0), zeroed[20]);
197 }
198
199 test "Highway lane tables gather across the full vector and two tables" {
200 const simd = @import("root.zig");
201 const D = simd.FixedTag(f32, 8);
202 const I = D.rebind(u32);
203 const first: D.Vector = .{ 0, 1, 2, 3, 4, 5, 6, 7 };
204 const second: D.Vector = .{ 8, 9, 10, 11, 12, 13, 14, 15 };
205 const indices: I.Vector = .{ 7, 0, 5, 2, 4, 1, 6, 3 };
206 try std.testing.expect(@reduce(.And, tableLookupLanes(D, first, indices) ==
207 @as(D.Vector, .{ 7, 0, 5, 2, 4, 1, 6, 3 })));
208 const both: I.Vector = .{ 15, 0, 9, 2, 12, 1, 14, 3 };
209 try std.testing.expect(@reduce(.And, twoTablesLookupLanes(D, first, second, both) ==
210 @as(D.Vector, .{ 15, 0, 9, 2, 12, 1, 14, 3 })));
211 }
212
213 test "Highway fixed table lookups and bit shuffle match scalar indexing" {
214 const simd = @import("root.zig");
215 const D = simd.FixedTag(u16, 8);
216 const I = D.rebind(u16);
217 const table = [_]u16{ 10, 11, 12, 13, 14, 15, 16, 17 };
218 const indices: I.Vector = .{ 7, 0, 6, 1, 5, 2, 4, 3 };
219 try std.testing.expect(@reduce(.And, lookup8(D, &table, indices) == @as(D.Vector, .{
220 17, 10, 16, 11, 15, 12, 14, 13,
221 })));
222
223 const B = simd.FixedTag(u64, 2);
224 const BI = B.repartition(u8);
225 const values: B.Vector = .{ 0x8000_0000_0000_0001, 0x0123_4567_89ab_cdef };
226 const bit_indices: BI.Vector = .{ 0, 63, 1, 62, 2, 61, 3, 60, 0, 4, 8, 12, 16, 20, 24, 28 };
227 try std.testing.expect(@reduce(.And, bitShuffle(B, values, bit_indices) == @as(B.Vector, .{ 3, 0x55 })));
228 }
229
230 fn verifyTableLaneType(comptime T: type) !void {
231 const simd = @import("root.zig");
232 const D = simd.FixedTag(T, 4);
233 const I = D.rebind(@Int(.unsigned, @bitSizeOf(T)));
234 const value: D.Vector = @splat(0);
235 const indices: I.Vector = .{ 3, 2, 1, 0 };
236 try std.testing.expect(@reduce(.And, tableLookupLanes(D, value, indicesFromVec(D, indices)) == value));
237 const source = [_]I.Lane{ 3, 2, 1, 0 };
238 try std.testing.expect(@reduce(.And, setTableIndices(D, &source) == indices));
239 }
240
241 test "Highway lane table setup instantiates every lane type" {
242 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {
243 try verifyTableLaneType(T);
244 }
245 }
246
247 test "Highway mixed byte tables and sixteen and thirty-two element gathers retain bits" {
248 const simd = @import("root.zig");
249 const Table = simd.FixedTag(u8, 16);
250 const D = simd.FixedTag(u32, 2);
251 const table = simd.iota(Table, 1);
252 const index_bytes = [8]u8{ 15, 0, 2, 1, 7, 6, 5, 4 };
253 const indices: D.Vector = @bitCast(index_bytes);
254 try std.testing.expectEqualSlices(u8, &.{ 16, 1, 3, 2, 8, 7, 6, 5 }, &@as([8]u8, @bitCast(tableLookupBytes(D, table, indices))));
255
256 const G = simd.FixedTag(u8, 16);
257 var table32: [32]u8 = undefined;
258 for (0..32) |index| table32[index] = @intCast(index + 10);
259 const indices16: G.Vector = .{ 15, 0, 14, 1, 13, 2, 12, 3, 11, 4, 10, 5, 9, 6, 8, 7 };
260 try std.testing.expectEqual(@as(u8, 25), lookup16(G, table32[0..16], indices16)[0]);
261 const indices32: G.Vector = .{ 31, 16, 30, 17, 29, 18, 28, 19, 27, 20, 26, 21, 25, 22, 24, 23 };
262 try std.testing.expectEqual(@as(u8, 41), lookup32(G, &table32, indices32)[0]);
263 try std.testing.expect(canLookup8(G) and canLookup16(G) and canLookup32(G));
264 }