lib/simd/src/compact.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn compress(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {
4 const lanes: [D.lane_count]D.Lane = @bitCast(value);
5 var result: [D.lane_count]D.Lane = undefined;
6 var position: usize = 0;
7 inline for (0..D.lane_count) |index| {
8 if (mask[index]) {
9 result[position] = lanes[index];
10 position += 1;
11 }
12 }
13 inline for (0..D.lane_count) |index| {
14 if (!mask[index]) {
15 result[position] = lanes[index];
16 position += 1;
17 }
18 }
19 std.debug.assert(position == D.lane_count);
20 return @bitCast(result);
21 }
22
23 pub fn compressIsPartition(comptime D: type) bool {
24 _ = D;
25 return true;
26 }
27
28 pub fn compressNot(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {
29 return compress(D, value, !mask);
30 }
31
32 pub fn compressBlocksNot(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {
33 if (D.Lane != u64) @compileError("compressBlocksNot requires u64 lanes");
34 if (D.lane_count < 2 or D.lane_count % 2 != 0) {
35 @compileError("compressBlocksNot requires complete 128-bit blocks");
36 }
37 return compressNot(D, value, mask);
38 }
39
40 pub fn compressStore(
41 comptime D: type,
42 value: D.Vector,
43 mask: D.Mask,
44 output: []D.Lane,
45 ) usize {
46 const count = countTrue(D, mask);
47 std.debug.assert(output.len >= count);
48 const compacted = compress(D, value, mask);
49 inline for (0..D.lane_count) |index| {
50 if (index < count) output[index] = compacted[index];
51 }
52 return count;
53 }
54
55 pub fn compressBlendedStore(
56 comptime D: type,
57 value: D.Vector,
58 mask: D.Mask,
59 output: []D.Lane,
60 ) usize {
61 return compressStore(D, value, mask, output);
62 }
63
64 pub fn compressBits(
65 comptime D: type,
66 value: D.Vector,
67 bits: []const u8,
68 ) D.Vector {
69 return compress(D, value, @import("compare.zig").loadMaskBits(D, bits));
70 }
71
72 pub fn compressBitsStore(
73 comptime D: type,
74 value: D.Vector,
75 bits: []const u8,
76 output: []D.Lane,
77 ) usize {
78 return compressStore(D, value, @import("compare.zig").loadMaskBits(D, bits), output);
79 }
80
81 pub fn expand(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {
82 const lanes: [D.lane_count]D.Lane = @bitCast(value);
83 var result: [D.lane_count]D.Lane = @splat(0);
84 var position: usize = 0;
85 inline for (0..D.lane_count) |index| {
86 if (mask[index]) {
87 result[index] = lanes[position];
88 position += 1;
89 }
90 }
91 return @bitCast(result);
92 }
93
94 pub fn loadExpand(
95 comptime D: type,
96 mask: D.Mask,
97 input: []const D.Lane,
98 ) D.Vector {
99 const count = countTrue(D, mask);
100 std.debug.assert(input.len >= count);
101 var result: D.Vector = @splat(0);
102 var position: usize = 0;
103 inline for (0..D.lane_count) |index| {
104 if (mask[index]) {
105 result[index] = input[position];
106 position += 1;
107 }
108 }
109 return result;
110 }
111
112 fn countTrue(comptime D: type, mask: D.Mask) usize {
113 return @import("compare.zig").countTrue(D, mask);
114 }
115
116 fn verifyLaneType(comptime T: type) !void {
117 const simd = @import("root.zig");
118 const D = simd.FixedTag(T, 8);
119 const value: D.Vector = @splat(0);
120 const mask: D.Mask = .{ true, false, true, false, false, true, false, true };
121 const compacted = compress(D, value, mask);
122 try std.testing.expect(@reduce(.And, expand(D, compacted, mask) == value));
123 var stored: [D.lane_count]T = undefined;
124 try std.testing.expectEqual(@as(usize, 4), compressStore(D, value, mask, &stored));
125 }
126
127 test "Highway compact operations instantiate every lane type" {
128 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {
129 try verifyLaneType(T);
130 }
131 try std.testing.expect(compressIsPartition(@import("root.zig").FixedTag(u32, 8)));
132 }
133
134 test "Highway compress is a stable partition and stores its selected prefix" {
135 const simd = @import("root.zig");
136 const D = simd.FixedTag(i32, 8);
137 const value: D.Vector = .{ 10, 11, 12, 13, 14, 15, 16, 17 };
138 const mask: D.Mask = .{ false, true, true, false, true, false, false, true };
139 const expected: D.Vector = .{ 11, 12, 14, 17, 10, 13, 15, 16 };
140 try std.testing.expect(@reduce(.And, compress(D, value, mask) == expected));
141 try std.testing.expect(@reduce(.And, compressNot(D, value, !mask) == expected));
142
143 var output = @as([8]i32, @splat(99));
144 try std.testing.expectEqual(@as(usize, 4), compressBlendedStore(D, value, mask, &output));
145 try std.testing.expectEqualSlices(i32, &.{ 11, 12, 14, 17, 99, 99, 99, 99 }, &output);
146
147 var bits = [_]u8{0x96};
148 try std.testing.expect(@reduce(.And, compressBits(D, value, &bits) == expected));
149 @memset(&output, 99);
150 try std.testing.expectEqual(@as(usize, 4), compressBitsStore(D, value, &bits, &output));
151 try std.testing.expectEqualSlices(i32, &.{ 11, 12, 14, 17, 99, 99, 99, 99 }, &output);
152 }
153
154 test "Highway expand scatters consecutive lanes and loadExpand reads only selected input" {
155 const simd = @import("root.zig");
156 const D = simd.FixedTag(u16, 8);
157 const compacted: D.Vector = .{ 20, 30, 50, 80, 91, 92, 93, 94 };
158 const mask: D.Mask = .{ false, true, true, false, true, false, false, true };
159 const expected: D.Vector = .{ 0, 20, 30, 0, 50, 0, 0, 80 };
160 try std.testing.expect(@reduce(.And, expand(D, compacted, mask) == expected));
161 try std.testing.expect(@reduce(.And, loadExpand(D, mask, &.{ 20, 30, 50, 80 }) == expected));
162 }
163
164 test "Highway block compression retains complete u64 pairs" {
165 const simd = @import("root.zig");
166 const D = simd.FixedTag(u64, 8);
167 const value: D.Vector = .{ 0, 1, 2, 3, 4, 5, 6, 7 };
168 const not_mask: D.Mask = .{ true, true, false, false, true, true, false, false };
169 try std.testing.expect(@reduce(.And, compressBlocksNot(D, value, not_mask) ==
170 @as(D.Vector, .{ 2, 3, 6, 7, 0, 1, 4, 5 })));
171 }