tiny.simd.tag

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

1const std = @import("std");23pub const maximum_vector_bytes: usize = 65_536;45pub fn FixedTag(comptime T: type, comptime lanes: usize) type {6    return Descriptor(T, lanes, 0);7}89pub fn Descriptor(comptime T: type, comptime lanes: usize, comptime power: i8) type {10    @setEvalBranchQuota(2_000);11    validateLane(T);12    if (power < -3) @compileError("Highway descriptor power must be at least -3");13    if (lanes == 0) @compileError("SIMD vectors require at least one lane");14    if (!std.math.isPowerOfTwo(lanes)) {15        @compileError("fixed SIMD lane counts must be powers of two");16    }17    if (lanes > maximum_vector_bytes / @sizeOf(T)) {18        @compileError("fixed SIMD vector exceeds the Highway maximum byte width");19    }20    const Storage = if (isBFloat16(T)) T.Storage else T;2122    return struct {23        pub const Lane = Storage;24        pub const Scalar = T;25        pub const lane_count: usize = lanes;26        pub const byte_count: usize = lanes * @sizeOf(T);27        pub const descriptor_power: i8 = power;28        pub const is_bfloat16 = isBFloat16(T);29        pub const Vector = @Vector(lanes, Storage);30        pub const Mask = @Vector(lanes, bool);3132        pub fn rebind(comptime NewLane: type) type {33            return Descriptor(34                NewLane,35                lanes,36                power + laneSizePower(NewLane) - laneSizePower(T),37            );38        }3940        pub fn repartition(comptime NewLane: type) type {41            validateLane(NewLane);42            const new_lanes = @max(1, (byte_count + @sizeOf(NewLane) - 1) /43                @sizeOf(NewLane));44            return Descriptor(NewLane, new_lanes, power);45        }4647        pub fn half() type {48            return Descriptor(T, @max(1, lanes / 2), power - 1);49        }5051        pub fn twice() type {52            return Descriptor(T, lanes * 2, power + 1);53        }5455        pub fn withLanes(comptime new_lanes: usize, comptime new_power: i8) type {56            return Descriptor(T, new_lanes, new_power);57        }58    };59}6061pub fn ScalableTag(comptime T: type) type {62    validateLane(T);63    return FixedTag(T, suggestedLanes(T));64}6566pub fn ScalableTagPow2(comptime T: type, comptime power: comptime_int) type {67    validateLane(T);68    const descriptor_power = clampedPower(power);69    return Descriptor(T, scaleLanes(suggestedLanes(T), descriptor_power), descriptor_power);70}7172pub fn CappedTag(comptime T: type, comptime limit: usize) type {73    validateLane(T);74    if (limit == 0) @compileError("capped SIMD vectors require a nonzero limit");75    const native_lanes = suggestedLanes(T);76    const lanes = std.math.floorPowerOfTwo(usize, @min(limit, native_lanes));77    return FixedTag(T, @max(1, lanes));78}7980pub fn CappedTagPow2(81    comptime T: type,82    comptime limit: usize,83    comptime power: comptime_int,84) type {85    const descriptor_power = clampedPower(power);86    return Descriptor(87        T,88        scaleLanes(CappedTag(T, limit).lane_count, descriptor_power),89        descriptor_power,90    );91}9293pub fn CappedTagIfFixed(comptime T: type, comptime limit: usize) type {94    return CappedTag(T, limit);95}9697pub fn CappedTagIfFixedPow2(98    comptime T: type,99    comptime limit: usize,100    comptime power: comptime_int,101) type {102    return CappedTagPow2(T, limit, power);103}104105pub fn Full16(comptime T: type) type {106    return FullBytes(T, 2);107}108109pub fn Full32(comptime T: type) type {110    return FullBytes(T, 4);111}112113pub fn Full64(comptime T: type) type {114    return FullBytes(T, 8);115}116117pub fn Full128(comptime T: type) type {118    return FullBytes(T, 16);119}120121pub fn TFromD(comptime D: type) type {122    return if (@hasDecl(D, "Scalar")) D.Scalar else D.Lane;123}124125pub fn DFromV(comptime V: type) type {126    return switch (@typeInfo(V)) {127        .vector => |info| FixedTag(info.child, info.len),128        else => @compileError("DFromV requires a Zig vector type"),129    };130}131132pub fn TFromV(comptime V: type) type {133    return switch (@typeInfo(V)) {134        .vector => |info| info.child,135        else => @compileError("TFromV requires a Zig vector type"),136    };137}138139pub fn VFromD(comptime D: type) type {140    return D.Vector;141}142143pub fn MFromD(comptime D: type) type {144    return D.Mask;145}146147pub fn Vec(comptime D: type) type {148    return VFromD(D);149}150151pub fn Mask(comptime D: type) type {152    return MFromD(D);153}154155pub fn maxLanes(comptime D: type) usize {156    return D.lane_count;157}158159pub fn laneCount(comptime D: type) usize {160    return D.lane_count;161}162163pub fn maxBytes(comptime D: type) usize {164    return D.byte_count;165}166167pub fn maxBlocks(comptime D: type) usize {168    return (D.byte_count + 15) / 16;169}170171pub fn pow2(comptime D: type) i8 {172    return D.descriptor_power;173}174175pub fn Rebind(comptime T: type, comptime D: type) type {176    return D.rebind(T);177}178179pub fn RebindToSigned(comptime D: type) type {180    return D.rebind(MakeSigned(TFromD(D)));181}182183pub fn RebindToUnsigned(comptime D: type) type {184    return D.rebind(MakeUnsigned(TFromD(D)));185}186187pub fn RebindToFloat(comptime D: type) type {188    return D.rebind(MakeFloat(TFromD(D)));189}190191pub fn Repartition(comptime T: type, comptime D: type) type {192    return D.repartition(T);193}194195pub fn RepartitionToWide(comptime D: type) type {196    return D.repartition(MakeWide(TFromD(D)));197}198199pub fn RepartitionToNarrow(comptime D: type) type {200    return D.repartition(MakeNarrow(TFromD(D)));201}202203pub fn RepartitionToWideX2(comptime D: type) type {204    return RepartitionToWide(RepartitionToWide(D));205}206207pub fn RepartitionToWideX3(comptime D: type) type {208    return RepartitionToWide(RepartitionToWideX2(D));209}210211pub fn Half(comptime D: type) type {212    return D.half();213}214215pub fn Twice(comptime D: type) type {216    return D.twice();217}218219pub fn BlockDFromD(comptime D: type) type {220    const block_lanes = @max(1, 16 / @sizeOf(TFromD(D)));221    return D.withLanes(@min(D.lane_count, block_lanes), 0);222}223224pub fn MakeUnsigned(comptime T: type) type {225    if (isBFloat16(T)) return u16;226    return switch (@typeInfo(T)) {227        .int => |info| switch (info.bits) {228            8, 16, 32, 64, 128 => @Int(.unsigned, info.bits),229            else => @compileError("unsupported Highway integer width"),230        },231        .float => |info| switch (info.bits) {232            16 => u16,233            32 => u32,234            64 => u64,235            else => @compileError("unsupported Highway floating-point width"),236        },237        else => @compileError("type has no Highway unsigned relation"),238    };239}240241pub fn MakeSigned(comptime T: type) type {242    if (isBFloat16(T)) return i16;243    return switch (@typeInfo(T)) {244        .int => |info| switch (info.bits) {245            8, 16, 32, 64 => @Int(.signed, info.bits),246            else => @compileError("type has no Highway signed relation"),247        },248        .float => |info| switch (info.bits) {249            16 => i16,250            32 => i32,251            64 => i64,252            else => @compileError("unsupported Highway floating-point width"),253        },254        else => @compileError("type has no Highway signed relation"),255    };256}257258pub fn MakeFloat(comptime T: type) type {259    if (isBFloat16(T)) @compileError("type has no Highway floating-point relation");260    return switch (@typeInfo(T)) {261        .int => |info| switch (info.bits) {262            16 => f16,263            32 => f32,264            64 => f64,265            else => @compileError("type has no Highway floating-point relation"),266        },267        .float => |info| switch (info.bits) {268            16 => f16,269            32 => f32,270            64 => f64,271            else => @compileError("type has no Highway floating-point relation"),272        },273        else => @compileError("type has no Highway floating-point relation"),274    };275}276277pub fn MakeWide(comptime T: type) type {278    if (isBFloat16(T)) return f32;279    return switch (@typeInfo(T)) {280        .int => |info| switch (info.bits) {281            8 => @Int(info.signedness, 16),282            16 => @Int(info.signedness, 32),283            32 => @Int(info.signedness, 64),284            64 => if (info.signedness == .unsigned)285                u128286            else287                @compileError("type has no Highway wide relation"),288            else => @compileError("type has no Highway wide relation"),289        },290        .float => |info| switch (info.bits) {291            16 => f32,292            32 => f64,293            else => @compileError("type has no Highway wide relation"),294        },295        else => @compileError("type has no Highway wide relation"),296    };297}298299pub fn MakeNarrow(comptime T: type) type {300    return switch (@typeInfo(T)) {301        .int => |info| switch (info.bits) {302            16 => @Int(info.signedness, 8),303            32 => @Int(info.signedness, 16),304            64 => @Int(info.signedness, 32),305            128 => if (info.signedness == .unsigned)306                u64307            else308                @compileError("type has no Highway narrow relation"),309            else => @compileError("type has no Highway narrow relation"),310        },311        .float => |info| switch (info.bits) {312            32 => f16,313            64 => f32,314            else => @compileError("type has no Highway narrow relation"),315        },316        else => @compileError("type has no Highway narrow relation"),317    };318}319320pub fn UnsignedFromSize(comptime bytes: usize) type {321    return switch (bytes) {322        1 => u8,323        2 => u16,324        4 => u32,325        8 => u64,326        16 => u128,327        else => @compileError("unsupported Highway unsigned byte width"),328    };329}330331pub fn SignedFromSize(comptime bytes: usize) type {332    return switch (bytes) {333        1 => i8,334        2 => i16,335        4 => i32,336        8 => i64,337        else => @compileError("unsupported Highway signed byte width"),338    };339}340341pub fn FloatFromSize(comptime bytes: usize) type {342    return switch (bytes) {343        2 => f16,344        4 => f32,345        8 => f64,346        else => @compileError("unsupported Highway floating-point byte width"),347    };348}349350pub fn isLane(comptime T: type) bool {351    if (isBFloat16(T)) return true;352    return switch (@typeInfo(T)) {353        .int => |info| info.bits == 8 or info.bits == 16 or354            info.bits == 32 or info.bits == 64,355        .float => |info| info.bits == 16 or info.bits == 32 or info.bits == 64,356        else => false,357    };358}359360pub fn isIntegerLane(comptime T: type) bool {361    return switch (@typeInfo(T)) {362        .int => |info| info.bits == 8 or info.bits == 16 or363            info.bits == 32 or info.bits == 64,364        else => false,365    };366}367368pub fn isSpecialFloat(comptime T: type) bool {369    return T == f16 or isBFloat16(T);370}371372pub fn isFloat3264(comptime T: type) bool {373    return T == f32 or T == f64;374}375376pub fn isFloat(comptime T: type) bool {377    return T == f16 or isFloat3264(T);378}379380pub fn isSigned(comptime T: type) bool {381    if (isBFloat16(T)) return true;382    return switch (@typeInfo(T)) {383        .int => |info| info.signedness == .signed,384        .float => true,385        else => false,386    };387}388389pub fn isUnsigned(comptime T: type) bool {390    return switch (@typeInfo(T)) {391        .int => |info| info.signedness == .unsigned,392        else => false,393    };394}395396fn FullBytes(comptime T: type, comptime bytes: usize) type {397    validateLane(T);398    if (@sizeOf(T) > bytes) @compileError("lane type exceeds fixed Highway vector width");399    return FixedTag(T, bytes / @sizeOf(T));400}401402fn scaleLanes(comptime lanes: usize, comptime power: i8) usize {403    if (power < 0) return @max(1, lanes >> @intCast(-power));404    return lanes << @intCast(power);405}406407fn clampedPower(comptime power: comptime_int) i8 {408    if (power < -3) @compileError("Highway scalable power must be at least -3");409    return @intCast(@min(power, 3));410}411412fn suggestedLanes(comptime T: type) usize {413    const Storage = if (isBFloat16(T)) T.Storage else T;414    return std.simd.suggestVectorLength(Storage) orelse 1;415}416417fn laneSizePower(comptime T: type) i8 {418    validateLane(T);419    return switch (@sizeOf(T)) {420        1 => 0,421        2 => 1,422        4 => 2,423        8 => 3,424        else => @compileError("unsupported Highway lane size"),425    };426}427428fn isBFloat16(comptime T: type) bool {429    return switch (@typeInfo(T)) {430        .@"struct" => @hasDecl(T, "is_bfloat16") and T.is_bfloat16,431        else => false,432    };433}434435fn validateLane(comptime T: type) void {436    if (!isLane(T)) {437        @compileError("SIMD lanes must be 8/16/32/64-bit integers or 16/32/64-bit floats");438    }439}440441test "fixed tags preserve lane and byte invariants" {442    const D = FixedTag(u16, 8);443    try std.testing.expectEqual(@as(usize, 8), D.lane_count);444    try std.testing.expectEqual(@as(usize, 16), D.byte_count);445    try std.testing.expectEqual(@as(i8, 0), D.descriptor_power);446    try std.testing.expectEqual(@Vector(8, u16), D.Vector);447    try std.testing.expectEqual(@Vector(8, bool), D.Mask);448}449450test "tag transformations match Highway descriptor meanings" {451    const D = FixedTag(u16, 8);452    try std.testing.expectEqual(@as(usize, 8), D.rebind(i32).lane_count);453    try std.testing.expectEqual(@as(i8, 1), D.rebind(i32).descriptor_power);454    try std.testing.expectEqual(FixedTag(u32, 4), D.repartition(u32));455    try std.testing.expectEqual(@as(usize, 4), D.half().lane_count);456    try std.testing.expectEqual(@as(i8, -1), D.half().descriptor_power);457    try std.testing.expectEqual(@as(usize, 16), D.twice().lane_count);458    try std.testing.expectEqual(@as(i8, 1), D.twice().descriptor_power);459}460461test "capped tags round down and full tags retain 128 bits" {462    const native = std.simd.suggestVectorLength(u8) orelse 1;463    const expected = std.math.floorPowerOfTwo(usize, @min(13, native));464    try std.testing.expectEqual(@max(1, expected), CappedTag(u8, 13).lane_count);465    try std.testing.expectEqual(@as(usize, 16), Full128(u8).lane_count);466    try std.testing.expectEqual(@as(usize, 2), Full128(f64).lane_count);467}468469test "Highway descriptor aliases preserve fixed widths powers and transforms" {470    const native = std.simd.suggestVectorLength(u32) orelse 1;471    try std.testing.expectEqual(native * 2, ScalableTagPow2(u32, 1).lane_count);472    try std.testing.expectEqual(@max(1, native / 2), ScalableTagPow2(u32, -1).lane_count);473    try std.testing.expectEqual(native * 8, ScalableTagPow2(u32, 99).lane_count);474    try std.testing.expectEqual(@as(i8, 3), pow2(ScalableTagPow2(u32, 99)));475    try std.testing.expectEqual(476        CappedTag(u32, 3).lane_count * 4,477        CappedTagPow2(u32, 3, 2).lane_count,478    );479    try std.testing.expectEqual(CappedTag(u32, 3), CappedTagIfFixed(u32, 3));480    try std.testing.expectEqual(@as(usize, 2), Full16(u8).lane_count);481    try std.testing.expectEqual(@as(usize, 2), Full32(u16).lane_count);482    try std.testing.expectEqual(@as(usize, 2), Full64(u32).lane_count);483484    const D = FixedTag(i16, 8);485    try std.testing.expectEqual(i16, TFromD(D));486    try std.testing.expectEqual(D, DFromV(D.Vector));487    try std.testing.expectEqual(i16, TFromV(D.Vector));488    try std.testing.expectEqual(D.Vector, VFromD(D));489    try std.testing.expectEqual(D.Mask, MFromD(D));490    try std.testing.expectEqual(D.Vector, Vec(D));491    try std.testing.expectEqual(D.Mask, Mask(D));492    try std.testing.expectEqual(D.lane_count, maxLanes(D));493    try std.testing.expectEqual(D.lane_count, laneCount(D));494    try std.testing.expectEqual(D.byte_count, maxBytes(D));495    try std.testing.expectEqual(@as(usize, 1), maxBlocks(D));496    try std.testing.expectEqual(@as(i8, 0), pow2(D));497    try std.testing.expectEqual(FixedTag(u16, 8), RebindToUnsigned(D));498    try std.testing.expectEqual(FixedTag(f16, 8), RebindToFloat(D));499    try std.testing.expectEqual(FixedTag(i32, 4), RepartitionToWide(D));500    try std.testing.expectEqual(FixedTag(i8, 16), RepartitionToNarrow(D));501    try std.testing.expectEqual(@as(usize, 4), Half(D).lane_count);502    try std.testing.expectEqual(@as(i8, -1), pow2(Half(D)));503    try std.testing.expectEqual(@as(usize, 16), Twice(D).lane_count);504    try std.testing.expectEqual(@as(i8, 1), pow2(Twice(D)));505    try std.testing.expectEqual(D, BlockDFromD(FixedTag(i16, 16)));506}507508test "Highway lane type relations cover every public scalar family" {509    try std.testing.expectEqual(u32, MakeUnsigned(f32));510    try std.testing.expectEqual(i64, MakeSigned(u64));511    try std.testing.expectEqual(f16, MakeFloat(u16));512    try std.testing.expectEqual(u128, MakeWide(u64));513    try std.testing.expectEqual(i16, MakeNarrow(i32));514    try std.testing.expectEqual(u128, UnsignedFromSize(16));515    try std.testing.expectEqual(i64, SignedFromSize(8));516    try std.testing.expectEqual(f64, FloatFromSize(8));517    try std.testing.expect(isIntegerLane(i8));518    try std.testing.expect(isSpecialFloat(f16));519    try std.testing.expect(isFloat3264(f64));520    try std.testing.expect(isFloat(f16));521    try std.testing.expect(isSigned(f32));522    try std.testing.expect(isUnsigned(u32));523}524525test "current Highway descriptor differential digest" {526    var digest: u64 = 0xcbf2_9ce4_8422_2325;527    const D = FixedTag(i16, 8);528    descriptorDigestRecord(D, &digest);529    descriptorDigestRecord(Rebind(i32, D), &digest);530    descriptorDigestRecord(Repartition(i32, D), &digest);531    descriptorDigestRecord(Half(D), &digest);532    descriptorDigestRecord(Twice(D), &digest);533    descriptorDigestRecord(BlockDFromD(Twice(D)), &digest);534    descriptorDigestRecord(Full16(u8), &digest);535    descriptorDigestRecord(Full32(u16), &digest);536    descriptorDigestRecord(Full64(u32), &digest);537    descriptorDigestRecord(Full128(u64), &digest);538    digest = descriptorDigestStep(539        digest,540        maxLanes(ScalableTagPow2(u32, 99)) / maxLanes(ScalableTag(u32)),541    );542    digest = descriptorDigestStep(digest, @bitCast(@as(i64, pow2(ScalableTagPow2(u32, 99)))));543    try std.testing.expectEqual(@as(u64, 9_371_378_310_310_814_615), digest);544}545546fn descriptorDigestRecord(comptime D: type, digest: *u64) void {547    digest.* = descriptorDigestStep(digest.*, maxLanes(D));548    digest.* = descriptorDigestStep(digest.*, maxBytes(D));549    digest.* = descriptorDigestStep(digest.*, maxBlocks(D));550    digest.* = descriptorDigestStep(digest.*, @bitCast(@as(i64, pow2(D))));551}552553fn descriptorDigestStep(digest: u64, value: u64) u64 {554    return (digest ^ value) *% 0x0000_0100_0000_01b3;555}