lib/simd/src/multiply.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const bfloat = @import("bfloat.zig");
3
4 pub fn mulRound(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
5 requireFloat(D.Lane, "mulRound");
6 var result: D.Vector = undefined;
7 inline for (0..D.lane_count) |index| {
8 result[index] = roundEvenScalar(D.Lane, a[index] * b[index]);
9 }
10 return result;
11 }
12
13 pub fn mulByPow2(comptime D: type, a: D.Vector, exponents: anytype) D.Vector {
14 requireFloat(D.Lane, "mulByPow2");
15 validateSameLanes(D, @TypeOf(exponents));
16 if (comptime @typeInfo(sourceLane(@TypeOf(exponents))) != .int) {
17 @compileError("mulByPow2 requires integer exponents");
18 }
19 var result: D.Vector = undefined;
20 inline for (0..D.lane_count) |index| {
21 result[index] = std.math.ldexp(a[index], exponentI32(exponents[index]));
22 }
23 return result;
24 }
25
26 pub fn mulByFloorPow2(comptime D: type, a: D.Vector, exponents: D.Vector) D.Vector {
27 requireFloat(D.Lane, "mulByFloorPow2");
28 var result: D.Vector = undefined;
29 inline for (0..D.lane_count) |index| {
30 const exponent = exponents[index];
31 if (std.math.isNan(exponent) or exponent == std.math.inf(D.Lane)) {
32 result[index] = a[index] * exponent;
33 } else if (exponent == -std.math.inf(D.Lane)) {
34 result[index] = a[index] * @as(D.Lane, 0);
35 } else {
36 result[index] = std.math.ldexp(a[index], floatExponentI32(@floor(exponent)));
37 }
38 }
39 return result;
40 }
41
42 pub fn mulHigh(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
43 requireInteger(D.Lane, "mulHigh");
44 const bits = @bitSizeOf(D.Lane);
45 const WideSigned = @Int(.signed, bits * 2);
46 const WideUnsigned = @Int(.unsigned, bits * 2);
47 const Unsigned = @Int(.unsigned, bits);
48 var result: D.Vector = undefined;
49 inline for (0..D.lane_count) |index| {
50 const high: Unsigned = if (@typeInfo(D.Lane).int.signedness == .signed) blk: {
51 const product = @as(WideSigned, a[index]) * @as(WideSigned, b[index]);
52 const product_bits: WideUnsigned = @bitCast(product);
53 break :blk @truncate(product_bits >> bits);
54 } else blk: {
55 const product = @as(WideUnsigned, a[index]) * @as(WideUnsigned, b[index]);
56 break :blk @truncate(product >> bits);
57 };
58 result[index] = @bitCast(high);
59 }
60 return result;
61 }
62
63 pub fn mulFixedPoint15(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
64 if (comptime D.Lane != i16) @compileError("mulFixedPoint15 requires i16 lanes");
65 var result: D.Vector = undefined;
66 inline for (0..D.lane_count) |index| {
67 const rounded = (@as(i32, a[index]) * @as(i32, b[index]) + 0x4000) >> 15;
68 result[index] = @intCast(std.math.clamp(rounded, std.math.minInt(i16), std.math.maxInt(i16)));
69 }
70 return result;
71 }
72
73 pub fn mulEven(comptime D: type, a: D.Vector, b: D.Vector) D.repartition(wideLane(D.Lane)).Vector {
74 return mulParity(D, a, b, 0);
75 }
76
77 pub fn mulOdd(comptime D: type, a: D.Vector, b: D.Vector) D.repartition(wideLane(D.Lane)).Vector {
78 return mulParity(D, a, b, 1);
79 }
80
81 pub fn widenMulPairwiseAdd(comptime DW: type, a: anytype, b: @TypeOf(a)) DW.Vector {
82 validatePairwiseWide(DW, @TypeOf(a));
83 var result: DW.Vector = undefined;
84 inline for (0..DW.lane_count) |index| {
85 const first = productLane(DW.Lane, a[index * 2], b[index * 2]);
86 const second = productLane(DW.Lane, a[index * 2 + 1], b[index * 2 + 1]);
87 result[index] = addLane(DW.Lane, first, second);
88 }
89 return result;
90 }
91
92 pub fn maskedWidenMulPairwiseAdd(
93 comptime DW: type,
94 mask: DW.Mask,
95 a: anytype,
96 b: @TypeOf(a),
97 ) DW.Vector {
98 return @select(
99 DW.Lane,
100 mask,
101 widenMulPairwiseAdd(DW, a, b),
102 @as(DW.Vector, @splat(0)),
103 );
104 }
105
106 pub fn satWidenMulPairwiseAdd(comptime DW: type, a: anytype, b: anytype) DW.Vector {
107 if (comptime DW.Lane != i16 or sourceLane(@TypeOf(a)) != u8 or sourceLane(@TypeOf(b)) != i8 or
108 sourceLaneCount(@TypeOf(a)) != DW.lane_count * 2 or
109 sourceLaneCount(@TypeOf(b)) != DW.lane_count * 2)
110 {
111 @compileError("satWidenMulPairwiseAdd requires paired u8 and i8 sources with i16 output");
112 }
113 var result: DW.Vector = undefined;
114 inline for (0..DW.lane_count) |index| {
115 const sum = @as(i32, a[index * 2]) * @as(i32, b[index * 2]) +
116 @as(i32, a[index * 2 + 1]) * @as(i32, b[index * 2 + 1]);
117 result[index] = @intCast(std.math.clamp(sum, std.math.minInt(i16), std.math.maxInt(i16)));
118 }
119 return result;
120 }
121
122 pub fn satWidenMulPairwiseAccumulate(
123 comptime DW: type,
124 a: anytype,
125 b: @TypeOf(a),
126 sum: DW.Vector,
127 ) DW.Vector {
128 if (comptime DW.Lane != i32 or sourceLane(@TypeOf(a)) != i16 or
129 sourceLaneCount(@TypeOf(a)) != DW.lane_count * 2)
130 {
131 @compileError("satWidenMulPairwiseAccumulate requires paired i16 sources with i32 output");
132 }
133 var result: DW.Vector = undefined;
134 inline for (0..DW.lane_count) |index| {
135 const value = @as(i64, sum[index]) +
136 @as(i64, a[index * 2]) * @as(i64, b[index * 2]) +
137 @as(i64, a[index * 2 + 1]) * @as(i64, b[index * 2 + 1]);
138 result[index] = @intCast(std.math.clamp(
139 value,
140 @as(i64, std.math.minInt(i32)),
141 @as(i64, std.math.maxInt(i32)),
142 ));
143 }
144 return result;
145 }
146
147 pub fn satWidenMulAccumFixedPoint(
148 comptime DW: type,
149 a: anytype,
150 b: @TypeOf(a),
151 sum: DW.Vector,
152 ) DW.Vector {
153 if (comptime DW.Lane != i32 or sourceLane(@TypeOf(a)) != i16 or
154 sourceLaneCount(@TypeOf(a)) != DW.lane_count)
155 {
156 @compileError("satWidenMulAccumFixedPoint requires same-lane-count i16 sources with i32 output");
157 }
158 var result: DW.Vector = undefined;
159 inline for (0..DW.lane_count) |index| {
160 const value = @as(i64, sum[index]) +
161 @as(i64, a[index]) * @as(i64, b[index]) * 2;
162 result[index] = @intCast(std.math.clamp(
163 value,
164 @as(i64, std.math.minInt(i32)),
165 @as(i64, std.math.maxInt(i32)),
166 ));
167 }
168 return result;
169 }
170
171 pub fn reorderWidenMulAccumulate(
172 comptime DW: type,
173 a: anytype,
174 b: @TypeOf(a),
175 sum0: DW.Vector,
176 sum1: *DW.Vector,
177 ) DW.Vector {
178 validatePairwiseWide(DW, @TypeOf(a));
179 var result = sum0;
180 inline for (0..DW.lane_count) |index| {
181 result[index] = addLane(DW.Lane, result[index], productLane(
182 DW.Lane,
183 a[index * 2],
184 b[index * 2],
185 ));
186 sum1[index] = addLane(DW.Lane, sum1[index], productLane(
187 DW.Lane,
188 a[index * 2 + 1],
189 b[index * 2 + 1],
190 ));
191 }
192 return result;
193 }
194
195 pub fn rearrangeToOddPlusEven(comptime D: type, sum0: D.Vector, sum1: D.Vector) D.Vector {
196 return if (@typeInfo(D.Lane) == .int) sum0 +% sum1 else sum0 + sum1;
197 }
198
199 pub fn sumOfMulQuadAccumulate(
200 comptime DW: type,
201 a: anytype,
202 b: anytype,
203 sum: DW.Vector,
204 ) DW.Vector {
205 if (comptime sourceLaneCount(@TypeOf(a)) != DW.lane_count * 4 or
206 sourceLaneCount(@TypeOf(b)) != DW.lane_count * 4)
207 {
208 @compileError("sumOfMulQuadAccumulate requires four source lanes per output lane");
209 }
210 var result = sum;
211 inline for (0..DW.lane_count) |index| {
212 inline for (0..4) |offset| {
213 result[index] = addLane(DW.Lane, result[index], productLane(
214 DW.Lane,
215 a[index * 4 + offset],
216 b[index * 4 + offset],
217 ));
218 }
219 }
220 return result;
221 }
222
223 pub fn widenMulAccumulate(
224 comptime DW: type,
225 a: anytype,
226 b: @TypeOf(a),
227 low: DW.Vector,
228 high: *DW.Vector,
229 ) DW.Vector {
230 validatePairwiseWide(DW, @TypeOf(a));
231 var result = low;
232 inline for (0..DW.lane_count) |index| {
233 result[index] = addLane(DW.Lane, result[index], productLane(
234 DW.Lane,
235 a[index],
236 b[index],
237 ));
238 high[index] = addLane(DW.Lane, high[index], productLane(
239 DW.Lane,
240 a[DW.lane_count + index],
241 b[DW.lane_count + index],
242 ));
243 }
244 return result;
245 }
246
247 pub fn maskedMulFixedPoint15(comptime D: type, mask: D.Mask, a: D.Vector, b: D.Vector) D.Vector {
248 return @select(D.Lane, mask, mulFixedPoint15(D, a, b), @as(D.Vector, @splat(0)));
249 }
250
251 fn mulParity(comptime D: type, a: D.Vector, b: D.Vector, comptime parity: usize) D.repartition(wideLane(D.Lane)).Vector {
252 requireInteger(D.Lane, "mulEven/mulOdd");
253 const Wide = wideLane(D.Lane);
254 const DW = D.repartition(Wide);
255 var result: DW.Vector = undefined;
256 if (@bitSizeOf(D.Lane) < 64) {
257 inline for (0..DW.lane_count) |index| {
258 result[index] = productLane(Wide, a[index * 2 + parity], b[index * 2 + parity]);
259 }
260 return result;
261 }
262 if (comptime D.lane_count & 1 != 0) @compileError("64-bit parity multiplication requires lane pairs");
263 const U128 = u128;
264 const I128 = i128;
265 inline for (0..D.lane_count / 2) |index| {
266 const source = index * 2 + parity;
267 const product_bits: U128 = if (@typeInfo(D.Lane).int.signedness == .signed)
268 @bitCast(@as(I128, a[source]) * @as(I128, b[source]))
269 else
270 @as(U128, a[source]) * @as(U128, b[source]);
271 const low: u64 = @truncate(product_bits);
272 const high: u64 = @truncate(product_bits >> 64);
273 result[index * 2] = @bitCast(low);
274 result[index * 2 + 1] = @bitCast(high);
275 }
276 return result;
277 }
278
279 fn productLane(comptime T: type, a: anytype, b: @TypeOf(a)) T {
280 if (comptime T == f32 and @TypeOf(a) == u16) {
281 return bfloat.f32FromBits(a) * bfloat.f32FromBits(b);
282 }
283 return switch (@typeInfo(T)) {
284 .int => @as(T, @intCast(a)) *% @as(T, @intCast(b)),
285 .float => @as(T, @floatCast(a)) * @as(T, @floatCast(b)),
286 else => unreachable,
287 };
288 }
289
290 fn addLane(comptime T: type, a: T, b: T) T {
291 return if (@typeInfo(T) == .int) a +% b else a + b;
292 }
293
294 fn wideLane(comptime T: type) type {
295 return switch (T) {
296 u8 => u16,
297 i8 => i16,
298 u16 => u32,
299 i16 => i32,
300 u32 => u64,
301 i32 => i64,
302 u64 => u64,
303 i64 => i64,
304 f16 => f32,
305 f32 => f64,
306 else => @compileError("lane has no Highway wide representation"),
307 };
308 }
309
310 fn roundEvenScalar(comptime T: type, value: T) T {
311 const rounded = @round(value);
312 const tie = @abs(value - @trunc(value)) == @as(T, 0.5);
313 const half = rounded / @as(T, 2);
314 if (tie and @trunc(half) != half) {
315 return rounded - if (value < 0) @as(T, -1) else @as(T, 1);
316 }
317 return rounded;
318 }
319
320 fn exponentI32(value: anytype) i32 {
321 const T = @TypeOf(value);
322 if (@typeInfo(T).int.signedness == .signed) {
323 if (value < std.math.minInt(i32)) return std.math.minInt(i32);
324 if (value > std.math.maxInt(i32)) return std.math.maxInt(i32);
325 return @intCast(value);
326 }
327 if (value > std.math.maxInt(i32)) return std.math.maxInt(i32);
328 return @intCast(value);
329 }
330
331 fn floatExponentI32(value: anytype) i32 {
332 if (value <= @as(@TypeOf(value), @floatFromInt(std.math.minInt(i32)))) return std.math.minInt(i32);
333 if (value >= @as(@TypeOf(value), @floatFromInt(std.math.maxInt(i32)))) return std.math.maxInt(i32);
334 return @intFromFloat(value);
335 }
336
337 fn validatePairwiseWide(comptime DW: type, comptime V: type) void {
338 if (comptime DW.Lane == f32 and sourceLane(V) == u16 and
339 sourceLaneCount(V) == DW.lane_count * 2)
340 {
341 return;
342 }
343 if (comptime sourceLaneCount(V) != DW.lane_count * 2 or wideLane(sourceLane(V)) != DW.Lane) {
344 @compileError("widening pair operation requires twice as many narrow source lanes");
345 }
346 }
347
348 fn validateSameLanes(comptime D: type, comptime V: type) void {
349 if (comptime sourceLaneCount(V) != D.lane_count) @compileError("vectors require equal lane counts");
350 }
351
352 fn sourceLane(comptime V: type) type {
353 return switch (@typeInfo(V)) {
354 .vector => |info| info.child,
355 else => @compileError("operation requires vector sources"),
356 };
357 }
358
359 fn sourceLaneCount(comptime V: type) usize {
360 return switch (@typeInfo(V)) {
361 .vector => |info| info.len,
362 else => @compileError("operation requires vector sources"),
363 };
364 }
365
366 fn requireInteger(comptime T: type, comptime operation: []const u8) void {
367 if (comptime @typeInfo(T) != .int) @compileError(operation ++ " requires integer lanes");
368 }
369
370 fn requireFloat(comptime T: type, comptime operation: []const u8) void {
371 if (comptime @typeInfo(T) != .float) @compileError(operation ++ " requires floating-point lanes");
372 }
373
374 test "Highway multiply variants preserve high halves parity and Q15 rounding" {
375 const simd = @import("root.zig");
376 const D = simd.FixedTag(i16, 8);
377 const a: D.Vector = .{ -32768, -3, -20000, 7, 16384, 9, 32767, -11 };
378 const b: D.Vector = .{ 32767, 5, -20000, 13, 16384, -17, 32767, 19 };
379 try std.testing.expect(@reduce(.And, mulHigh(D, a, b) == @as(D.Vector, .{
380 -16384, -1, 6103, 0, 4096, -1, 16383, -1,
381 })));
382 try std.testing.expect(@reduce(.And, mulFixedPoint15(D, a, b) == @as(D.Vector, .{
383 -32767, 0, 12207, 0, 8192, 0, 32766, 0,
384 })));
385 try std.testing.expect(@reduce(.And, mulEven(D, a, b) == @as(simd.FixedTag(i32, 4).Vector, .{
386 -1073709056, 400000000, 268435456, 1073676289,
387 })));
388 try std.testing.expect(@reduce(.And, mulOdd(D, a, b) == @as(simd.FixedTag(i32, 4).Vector, .{
389 -15, 91, -153, -209,
390 })));
391 }
392
393 test "Highway rounded and power-of-two multiplication follow floating semantics" {
394 const simd = @import("root.zig");
395 const D = simd.FixedTag(f32, 4);
396 const E = simd.FixedTag(i32, 4);
397 const a: D.Vector = .{ -3.5, -2.5, 2.5, 3.5 };
398 try std.testing.expect(@reduce(.And, mulRound(D, a, @as(D.Vector, @splat(1))) ==
399 @as(D.Vector, .{ -4, -2, 2, 4 })));
400 try std.testing.expect(@reduce(.And, mulByPow2(D, @as(D.Vector, @splat(1.5)), @as(E.Vector, .{ -2, -1, 1, 2 })) ==
401 @as(D.Vector, .{ 0.375, 0.75, 3, 6 })));
402 try std.testing.expect(@reduce(.And, mulByFloorPow2(D, @as(D.Vector, @splat(2)), @as(D.Vector, .{ -1.2, -0.2, 1.8, 2.9 })) ==
403 @as(D.Vector, .{ 0.5, 1, 4, 8 })));
404 const exceptional = mulByFloorPow2(D, @as(D.Vector, @splat(2)), @as(D.Vector, .{
405 -std.math.inf(f32), std.math.inf(f32), std.math.nan(f32), 0,
406 }));
407 try std.testing.expectEqual(@as(f32, 0), exceptional[0]);
408 try std.testing.expectEqual(std.math.inf(f32), exceptional[1]);
409 try std.testing.expect(std.math.isNan(exceptional[2]));
410 }
411
412 test "Highway widening multiply families fold consecutive source groups" {
413 const simd = @import("root.zig");
414 const N = simd.FixedTag(i16, 8);
415 const W = simd.FixedTag(i32, 4);
416 const a: N.Vector = .{ 1, 2, 3, 4, 5, 6, 7, 8 };
417 const b: N.Vector = .{ 8, 7, 6, 5, 4, 3, 2, 1 };
418 try std.testing.expect(@reduce(.And, widenMulPairwiseAdd(W, a, b) ==
419 @as(W.Vector, .{ 22, 38, 38, 22 })));
420 var odd: W.Vector = @splat(0);
421 const even = reorderWidenMulAccumulate(W, a, b, @splat(0), &odd);
422 try std.testing.expect(@reduce(.And, rearrangeToOddPlusEven(W, even, odd) ==
423 @as(W.Vector, .{ 22, 38, 38, 22 })));
424 var high: W.Vector = @splat(10);
425 const low = widenMulAccumulate(W, a, b, @splat(1), &high);
426 try std.testing.expect(@reduce(.And, low == @as(W.Vector, .{ 9, 15, 19, 21 })));
427 try std.testing.expect(@reduce(.And, high == @as(W.Vector, .{ 30, 28, 24, 18 })));
428 const Q = simd.FixedTag(i64, 2);
429 try std.testing.expect(@reduce(.And, sumOfMulQuadAccumulate(Q, a, b, @as(Q.Vector, @splat(3))) ==
430 @as(Q.Vector, .{ 63, 63 })));
431 }
432
433 test "Highway saturating widening multiply handles mixed signedness and accumulation" {
434 const simd = @import("root.zig");
435 const U = simd.FixedTag(u8, 8);
436 const I = simd.FixedTag(i8, 8);
437 const W = simd.FixedTag(i16, 4);
438 const a: U.Vector = .{ 255, 255, 1, 2, 100, 200, 255, 255 };
439 const b: I.Vector = .{ 127, 127, -3, 4, -100, 100, -128, -128 };
440 try std.testing.expect(@reduce(.And, satWidenMulPairwiseAdd(W, a, b) ==
441 @as(W.Vector, .{ 32767, 5, 10000, -32768 })));
442 const N = simd.FixedTag(i16, 8);
443 const DW = simd.FixedTag(i32, 4);
444 const x: N.Vector = @splat(-32768);
445 try std.testing.expect(@reduce(.And, satWidenMulPairwiseAccumulate(DW, x, x, @splat(1)) ==
446 @as(DW.Vector, @splat(std.math.maxInt(i32)))));
447 const Same = simd.FixedTag(i16, 4);
448 try std.testing.expect(@reduce(.And, satWidenMulAccumFixedPoint(DW, @as(Same.Vector, @splat(-32768)), @as(Same.Vector, @splat(-32768)), @splat(0)) ==
449 @as(DW.Vector, @splat(std.math.maxInt(i32)))));
450 }