lib/simd/src/count.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 pub 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 }
10 
11 pub 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 }
18 
19 pub 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 }
26 
27 pub 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 }
34 
35 pub 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 }
43 
44 fn UnsignedVector(comptime D: type) type {
45     return @Vector(D.lane_count, @Int(.unsigned, @bitSizeOf(D.Lane)));
46 }
47 
48 fn validate(comptime D: type) void {
49     if (@typeInfo(D.Lane) != .int) @compileError("bit counts require integer lanes");
50 }
51 
52 fn 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 }
69 
70 test "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 }
75 
76 test "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 }