tiny.simd.table

1pub const table = @import("table.zig");

1const std = @import("std");23pub fn tableLookupBytes(comptime D: type, table: anytype, indices: D.Vector) D.Vector {4    return lookupBytes(D, false, table, indices);5}67pub fn tableLookupBytesOr0(comptime D: type, table: anytype, indices: D.Vector) D.Vector {8    return lookupBytes(D, true, table, indices);9}1011pub 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}3132pub 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 or35        @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}4142pub 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}4950pub 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}6263pub 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        else80            second_lanes[index - D.lane_count];81    }82    return @bitCast(result);83}8485pub fn lookup8(comptime D: type, table: []const D.Lane, indices: anytype) D.Vector {86    return lookupN(D, 8, table, indices);87}8889pub fn lookup16(comptime D: type, table: []const D.Lane, indices: anytype) D.Vector {90    return lookupN(D, 16, table, indices);91}9293pub fn lookup32(comptime D: type, table: []const D.Lane, indices: anytype) D.Vector {94    return lookupN(D, 32, table, indices);95}9697pub fn canLookup8(comptime D: type) bool {98    _ = D;99    return true;100}101102pub fn canLookup16(comptime D: type) bool {103    _ = D;104    return true;105}106107pub fn canLookup32(comptime D: type) bool {108    _ = D;109    return true;110}111112fn 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}133134fn 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}151152fn 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 or155        @bitSizeOf(info.child) != @bitSizeOf(D.Lane))156    {157        @compileError("lane lookup indices require enough same-width integer lanes");158    }159    return info;160}161162fn 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}168169fn 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}174175fn indexLane(comptime D: type) type {176    return @Int(.unsigned, @bitSizeOf(D.Lane));177}178179test "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}198199test "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}212213test "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    })));222223    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}229230fn 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}240241test "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}246247test "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))));255256    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}