tiny.simd.count

1const std = @import("std");23pub fn populationCount(comptime D: type, value: D.Vector) D.Vector {4    validate(D);5    if (@typeInfo(D.Lane).int.signedness != .unsigned) {6        @compileError("populationCount requires unsigned lanes");7    }8    return @popCount(value);9}1011pub fn leadingZeroCount(comptime D: type, value: D.Vector) D.Vector {12    validate(D);13    const U = UnsignedVector(D);14    const bits: U = @bitCast(value);15    const result: U = @intCast(@clz(bits));16    return @bitCast(result);17}1819pub fn maskedLeadingZeroCount(20    comptime D: type,21    mask: D.Mask,22    value: D.Vector,23) D.Vector {24    return @select(D.Lane, mask, leadingZeroCount(D, value), @as(D.Vector, @splat(0)));25}2627pub fn trailingZeroCount(comptime D: type, value: D.Vector) D.Vector {28    validate(D);29    const U = UnsignedVector(D);30    const bits: U = @bitCast(value);31    const result: U = @intCast(@ctz(bits));32    return @bitCast(result);33}3435pub fn highestSetBitIndex(comptime D: type, value: D.Vector) D.Vector {36    validate(D);37    const U = UnsignedVector(D);38    const bits: U = @bitCast(value);39    const top: U = @splat(@bitSizeOf(D.Lane) - 1);40    const leading: U = @intCast(@clz(bits));41    return @bitCast(top -% leading);42}4344fn UnsignedVector(comptime D: type) type {45    return @Vector(D.lane_count, @Int(.unsigned, @bitSizeOf(D.Lane)));46}4748fn validate(comptime D: type) void {49    if (@typeInfo(D.Lane) != .int) @compileError("bit counts require integer lanes");50}5152fn verifyCounts(comptime T: type) !void {53    const simd = @import("root.zig");54    const D = simd.FixedTag(T, 4);55    const U = @Int(.unsigned, @bitSizeOf(T));56    const UV = @Vector(4, U);57    const bits = @bitSizeOf(T);58    const input_bits: UV = .{ 0, 1, @as(U, 1) << (bits - 1), 0x2b };59    const input: D.Vector = @bitCast(input_bits);60    const leading: D.Vector = @bitCast(@as(UV, .{ bits, bits - 1, 0, bits - 6 }));61    const trailing: D.Vector = @bitCast(@as(UV, .{ bits, 0, bits - 1, 0 }));62    const highest: D.Vector = @bitCast(@as(UV, .{63        std.math.maxInt(U), 0, bits - 1, 5,64    }));65    try std.testing.expect(@reduce(.And, leadingZeroCount(D, input) == leading));66    try std.testing.expect(@reduce(.And, trailingZeroCount(D, input) == trailing));67    try std.testing.expect(@reduce(.And, highestSetBitIndex(D, input) == highest));68}6970test "Highway bit counts cover every integer lane width" {71    inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64 }) |T| {72        try verifyCounts(T);73    }74}7576test "Highway population and masked leading counts retain lane results" {77    const simd = @import("root.zig");78    const D = simd.FixedTag(u32, 4);79    const value: D.Vector = .{ 0, 1, 0xffff_ffff, 0xf0f0_0001 };80    const mask: D.Mask = .{ true, false, true, false };81    try std.testing.expect(@reduce(.And, populationCount(D, value) ==82        @as(D.Vector, .{ 0, 1, 32, 9 })));83    try std.testing.expect(@reduce(.And, maskedLeadingZeroCount(D, mask, value) ==84        @as(D.Vector, .{ 32, 0, 0, 0 })));85}

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