lib/simd/src/reduce.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn sum(comptime D: type, value: D.Vector) D.Lane {
4 var result: D.Lane = 0;
5 inline for (0..D.lane_count) |index| {
6 result = if (comptime @typeInfo(D.Lane) == .int)
7 result +% value[index]
8 else
9 result + value[index];
10 }
11 return result;
12 }
13
14 pub fn min(comptime D: type, value: D.Vector) D.Lane {
15 var result = value[0];
16 inline for (1..D.lane_count) |index| {
17 result = @min(result, value[index]);
18 }
19 return result;
20 }
21
22 pub fn max(comptime D: type, value: D.Vector) D.Lane {
23 var result = value[0];
24 inline for (1..D.lane_count) |index| {
25 result = @max(result, value[index]);
26 }
27 return result;
28 }
29
30 pub fn sumOfLanes(comptime D: type, value: D.Vector) D.Vector {
31 return @splat(sum(D, value));
32 }
33
34 pub fn minOfLanes(comptime D: type, value: D.Vector) D.Vector {
35 return @splat(min(D, value));
36 }
37
38 pub fn maxOfLanes(comptime D: type, value: D.Vector) D.Vector {
39 return @splat(max(D, value));
40 }
41
42 pub fn maskedSum(comptime D: type, mask: D.Mask, value: D.Vector) D.Lane {
43 var result: D.Lane = 0;
44 inline for (0..D.lane_count) |index| {
45 if (mask[index]) {
46 result = if (comptime @typeInfo(D.Lane) == .int)
47 result +% value[index]
48 else
49 result + value[index];
50 }
51 }
52 return result;
53 }
54
55 pub fn maskedMin(comptime D: type, mask: D.Mask, value: D.Vector) D.Lane {
56 var result: D.Lane = 0;
57 var found = false;
58 inline for (0..D.lane_count) |index| {
59 if (mask[index]) {
60 result = if (found) @min(result, value[index]) else value[index];
61 found = true;
62 }
63 }
64 return result;
65 }
66
67 pub fn maskedMax(comptime D: type, mask: D.Mask, value: D.Vector) D.Lane {
68 var result: D.Lane = 0;
69 var found = false;
70 inline for (0..D.lane_count) |index| {
71 if (mask[index]) {
72 result = if (found) @max(result, value[index]) else value[index];
73 found = true;
74 }
75 }
76 return result;
77 }
78
79 pub fn sumsOf2(comptime D: type, value: D.Vector) D.repartition(wideLane(D.Lane)).Vector {
80 const W = wideLane(D.Lane);
81 const R = D.repartition(W);
82 if (D.lane_count < 2) @compileError("sumsOf2 requires at least two lanes");
83 var result: R.Vector = undefined;
84 inline for (0..R.lane_count) |index| {
85 result[index] = widen(W, value[index * 2]) + widen(W, value[index * 2 + 1]);
86 }
87 return result;
88 }
89
90 pub fn sumsOf4(comptime D: type, value: D.Vector) D.repartition(wideLane(wideLane(D.Lane))).Vector {
91 const W = wideLane(wideLane(D.Lane));
92 const R = D.repartition(W);
93 if (D.lane_count < 4) @compileError("sumsOf4 requires at least four lanes");
94 var result: R.Vector = undefined;
95 inline for (0..R.lane_count) |index| {
96 const base = index * 4;
97 result[index] = widen(W, value[base]) + widen(W, value[base + 1]) +
98 widen(W, value[base + 2]) + widen(W, value[base + 3]);
99 }
100 return result;
101 }
102
103 pub fn sumsOf8(comptime D: type, value: D.Vector) D.repartition(wideLane(wideLane(wideLane(D.Lane)))).Vector {
104 const W = wideLane(wideLane(wideLane(D.Lane)));
105 const R = D.repartition(W);
106 if (D.lane_count < 8) @compileError("sumsOf8 requires at least eight lanes");
107 var result: R.Vector = undefined;
108 inline for (0..R.lane_count) |index| {
109 const base = index * 8;
110 result[index] = widen(W, value[base]) + widen(W, value[base + 1]) +
111 widen(W, value[base + 2]) + widen(W, value[base + 3]) +
112 widen(W, value[base + 4]) + widen(W, value[base + 5]) +
113 widen(W, value[base + 6]) + widen(W, value[base + 7]);
114 }
115 return result;
116 }
117
118 pub fn sumsOf8AbsDiff(
119 comptime D: type,
120 a: D.Vector,
121 b: D.Vector,
122 ) D.repartition(wideLane(wideLane(wideLane(D.Lane)))).Vector {
123 requireByteLane(D, "sumsOf8AbsDiff");
124 const W = wideLane(wideLane(wideLane(D.Lane)));
125 const R = D.repartition(W);
126 if (D.lane_count < 8) @compileError("sumsOf8AbsDiff requires at least eight lanes");
127 var result: R.Vector = undefined;
128 inline for (0..R.lane_count) |index| {
129 var total: W = 0;
130 inline for (0..8) |offset| {
131 total += absoluteDifference(W, a[index * 8 + offset], b[index * 8 + offset]);
132 }
133 result[index] = total;
134 }
135 return result;
136 }
137
138 pub fn sumsOfAdjQuadAbsDiff(
139 comptime D: type,
140 comptime a_offset: usize,
141 comptime b_offset: usize,
142 a: D.Vector,
143 b: D.Vector,
144 ) D.repartition(wideLane(D.Lane)).Vector {
145 requireByteLane(D, "sumsOfAdjQuadAbsDiff");
146 if (a_offset > 1 or b_offset > 3) @compileError("quad offsets exceed a 128-bit block");
147 const W = wideLane(D.Lane);
148 const R = D.repartition(W);
149 var result: R.Vector = @splat(0);
150 inline for (0..R.lane_count) |index| {
151 const a_base = a_offset * 4 + (index / 8) * 16 + (index & 7);
152 const b_base = b_offset * 4 + (index / 8) * 16;
153 if (a_base + 3 < D.lane_count and b_base + 3 < D.lane_count) {
154 var total: W = 0;
155 inline for (0..4) |offset| {
156 total += absoluteDifference(W, a[a_base + offset], b[b_base + offset]);
157 }
158 result[index] = total;
159 }
160 }
161 return result;
162 }
163
164 pub fn sumsOfShuffledQuadAbsDiff(
165 comptime D: type,
166 comptime index3: usize,
167 comptime index2: usize,
168 comptime index1: usize,
169 comptime index0: usize,
170 a: D.Vector,
171 b: D.Vector,
172 ) D.repartition(wideLane(D.Lane)).Vector {
173 requireByteLane(D, "sumsOfShuffledQuadAbsDiff");
174 if (index0 > 3 or index1 > 3 or index2 > 3 or index3 > 3) {
175 @compileError("quad shuffle indices must be less than four");
176 }
177 const W = wideLane(D.Lane);
178 const R = D.repartition(W);
179 const selectors = [4]usize{ index0, index1, index2, index3 };
180 var shuffled: [D.lane_count]D.Lane = @splat(0);
181 inline for (0..(D.lane_count + 15) / 16) |block_index| {
182 inline for (0..4) |group| {
183 inline for (0..4) |byte| {
184 const destination = block_index * 16 + group * 4 + byte;
185 const source = block_index * 16 + selectors[group] * 4 + byte;
186 if (destination < D.lane_count and source < D.lane_count) {
187 shuffled[destination] = a[source];
188 }
189 }
190 }
191 }
192 var result: R.Vector = @splat(0);
193 inline for (0..R.lane_count) |index| {
194 const a_base = (index / 4) * 8 + (index & 3);
195 const b_base = (index / 2) * 4;
196 if (a_base + 3 < D.lane_count and b_base + 3 < D.lane_count) {
197 var total: W = 0;
198 inline for (0..4) |offset| {
199 total += absoluteDifference(W, shuffled[a_base + offset], b[b_base + offset]);
200 }
201 result[index] = total;
202 }
203 }
204 return result;
205 }
206
207 fn wideLane(comptime T: type) type {
208 return switch (T) {
209 i8 => i16,
210 u8 => u16,
211 i16 => i32,
212 u16 => u32,
213 i32 => i64,
214 u32 => u64,
215 f16 => f32,
216 f32 => f64,
217 else => @compileError("lane type has no Highway wide representation"),
218 };
219 }
220
221 fn widen(comptime T: type, value: anytype) T {
222 return if (@typeInfo(T) == .float) @floatCast(value) else @intCast(value);
223 }
224
225 fn absoluteDifference(comptime T: type, a: anytype, b: @TypeOf(a)) T {
226 const a_wide: i128 = a;
227 const b_wide: i128 = b;
228 const difference = a_wide - b_wide;
229 return @intCast(if (difference < 0) -difference else difference);
230 }
231
232 fn requireByteLane(comptime D: type, comptime operation: []const u8) void {
233 if (comptime D.Lane != u8 and D.Lane != i8) {
234 @compileError(operation ++ " requires byte integer lanes");
235 }
236 }
237
238 test "reductions match scalar Highway models" {
239 const simd = @import("root.zig");
240 const D = simd.FixedTag(i16, 8);
241 const value: D.Vector = .{ 7, -4, 2, 9, -11, 3, 1, 6 };
242 try std.testing.expectEqual(@as(i16, 13), sum(D, value));
243 try std.testing.expectEqual(@as(i16, -11), min(D, value));
244 try std.testing.expectEqual(@as(i16, 9), max(D, value));
245 }
246
247 test "integer sum uses Highway wrapping semantics" {
248 const simd = @import("root.zig");
249 const D = simd.FixedTag(u8, 4);
250 try std.testing.expectEqual(@as(u8, 4), sum(D, @as(D.Vector, @splat(129))));
251 }
252
253 test "Highway broadcast and masked reductions match scalar selection" {
254 const simd = @import("root.zig");
255 const D = simd.FixedTag(i16, 8);
256 const value: D.Vector = .{ 7, -4, 2, 9, -11, 3, 1, 6 };
257 const mask: D.Mask = .{ false, true, false, true, true, false, true, false };
258 try std.testing.expect(@reduce(.And, sumOfLanes(D, value) == @as(D.Vector, @splat(13))));
259 try std.testing.expect(@reduce(.And, minOfLanes(D, value) == @as(D.Vector, @splat(-11))));
260 try std.testing.expect(@reduce(.And, maxOfLanes(D, value) == @as(D.Vector, @splat(9))));
261 try std.testing.expectEqual(@as(i16, -5), maskedSum(D, mask, value));
262 try std.testing.expectEqual(@as(i16, -11), maskedMin(D, mask, value));
263 try std.testing.expectEqual(@as(i16, 9), maskedMax(D, mask, value));
264 }
265
266 test "Highway pairwise sums widen adjacent groups" {
267 const simd = @import("root.zig");
268 const D8 = simd.FixedTag(i8, 16);
269 const value: D8.Vector = .{
270 100, 100, -100, -100, 1, 2, 3, 4,
271 5, 6, 7, 8, 9, 10, 11, 12,
272 };
273 const D16 = D8.repartition(i16);
274 const D32 = D8.repartition(i32);
275 const D64 = D8.repartition(i64);
276 try std.testing.expect(@reduce(.And, sumsOf2(D8, value) == @as(D16.Vector, .{
277 200, -200, 3, 7, 11, 15, 19, 23,
278 })));
279 try std.testing.expect(@reduce(.And, sumsOf4(D8, value) == @as(D32.Vector, .{ 0, 10, 26, 42 })));
280 try std.testing.expect(@reduce(.And, sumsOf8(D8, value) == @as(D64.Vector, .{ 10, 68 })));
281
282 const F = simd.FixedTag(f32, 4);
283 try std.testing.expect(@reduce(.And, sumsOf2(F, @as(F.Vector, .{ 1.5, 2.5, -4, 1 })) ==
284 @as(F.repartition(f64).Vector, .{ 4, -3 })));
285 }
286
287 fn verifyReductionLaneType(comptime T: type) !void {
288 const simd = @import("root.zig");
289 const D = simd.FixedTag(T, 4);
290 const value: D.Vector = @splat(0);
291 const mask: D.Mask = .{ true, false, true, false };
292 try std.testing.expectEqual(@as(T, 0), sum(D, value));
293 try std.testing.expectEqual(@as(T, 0), min(D, value));
294 try std.testing.expectEqual(@as(T, 0), max(D, value));
295 try std.testing.expectEqual(@as(T, 0), maskedSum(D, mask, value));
296 try std.testing.expectEqual(@as(T, 0), maskedMin(D, mask, value));
297 try std.testing.expectEqual(@as(T, 0), maskedMax(D, mask, value));
298 try std.testing.expect(@reduce(.And, sumOfLanes(D, value) == value));
299 try std.testing.expect(@reduce(.And, minOfLanes(D, value) == value));
300 try std.testing.expect(@reduce(.And, maxOfLanes(D, value) == value));
301 }
302
303 fn verifySumsOf2(comptime T: type) !void {
304 const simd = @import("root.zig");
305 const D = simd.FixedTag(T, 8);
306 const value: D.Vector = @splat(0);
307 try std.testing.expect(@reduce(.And, sumsOf2(D, value) == @as(D.repartition(wideLane(T)).Vector, @splat(0))));
308 }
309
310 fn verifySumsOf4(comptime T: type) !void {
311 const simd = @import("root.zig");
312 const D = simd.FixedTag(T, 8);
313 const W = wideLane(wideLane(T));
314 const value: D.Vector = @splat(0);
315 try std.testing.expect(@reduce(.And, sumsOf4(D, value) == @as(D.repartition(W).Vector, @splat(0))));
316 }
317
318 test "Highway reductions instantiate every lane type and supported widening" {
319 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {
320 try verifyReductionLaneType(T);
321 }
322 inline for (.{ u8, i8, u16, i16, u32, i32, f16, f32 }) |T| try verifySumsOf2(T);
323 inline for (.{ u8, i8, u16, i16 }) |T| try verifySumsOf4(T);
324 const simd = @import("root.zig");
325 inline for (.{ u8, i8 }) |T| {
326 const D = simd.FixedTag(T, 8);
327 try std.testing.expect(@reduce(.And, sumsOf8(D, @as(D.Vector, @splat(0))) ==
328 @as(D.repartition(wideLane(wideLane(wideLane(T)))).Vector, @splat(0))));
329 }
330 }
331
332 test "Highway absolute-difference reductions match scalar block formulas" {
333 const simd = @import("root.zig");
334 const D = simd.FixedTag(i8, 16);
335 const a: D.Vector = .{ 1, 2, 3, 4, 9, 8, 7, 6, -1, -2, -3, -4, 5, 6, 7, 8 };
336 const b: D.Vector = .{ 4, 3, 2, 1, 1, 2, 3, 4, 1, 2, 3, 4, 8, 7, 6, 5 };
337 try std.testing.expect(@reduce(.And, sumsOf8AbsDiff(D, a, b) ==
338 @as(D.repartition(i64).Vector, .{ 28, 28 })));
339 try std.testing.expect(@reduce(.And, sumsOfAdjQuadAbsDiff(D, 0, 1, a, b) ==
340 @as(D.repartition(i16).Vector, .{ 0, 8, 14, 18, 20, 20, 20, 20 })));
341 try std.testing.expect(@reduce(.And, sumsOfShuffledQuadAbsDiff(D, 3, 2, 1, 0, a, b) ==
342 @as(D.repartition(i16).Vector, .{ 8, 12, 14, 18, 20, 16, 24, 16 })));
343 }