lib/simd/src/math.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const Unary = enum {
4 acos,
5 acosh,
6 asin,
7 asinh,
8 atan,
9 atanh,
10 cbrt,
11 expm1,
12 log10,
13 log1p,
14 sinh,
15 cosh,
16 tanh,
17 tgamma,
18 log_gamma,
19 };
20
21 const Binary = enum {
22 atan2,
23 hypot,
24 pow,
25 };
26
27 pub fn acos(comptime D: type, value: D.Vector) D.Vector {
28 return mapUnary(D, value, .acos);
29 }
30
31 pub fn acosh(comptime D: type, value: D.Vector) D.Vector {
32 return mapUnary(D, value, .acosh);
33 }
34
35 pub fn asin(comptime D: type, value: D.Vector) D.Vector {
36 return mapUnary(D, value, .asin);
37 }
38
39 pub fn asinh(comptime D: type, value: D.Vector) D.Vector {
40 return mapUnary(D, value, .asinh);
41 }
42
43 pub fn atan(comptime D: type, value: D.Vector) D.Vector {
44 return mapUnary(D, value, .atan);
45 }
46
47 pub fn atanh(comptime D: type, value: D.Vector) D.Vector {
48 return mapUnary(D, value, .atanh);
49 }
50
51 pub fn atan2(comptime D: type, y: D.Vector, x: D.Vector) D.Vector {
52 return mapBinary(D, y, x, .atan2);
53 }
54
55 pub fn cbrt(comptime D: type, value: D.Vector) D.Vector {
56 return mapUnary(D, value, .cbrt);
57 }
58
59 pub fn cbrtNormal(comptime D: type, value: D.Vector) D.Vector {
60 return mapUnary(D, value, .cbrt);
61 }
62
63 pub fn cos(comptime D: type, value: D.Vector) D.Vector {
64 requireFloat(D.Lane);
65 return @cos(value);
66 }
67
68 pub fn tan(comptime D: type, value: D.Vector) D.Vector {
69 requireFloat(D.Lane);
70 return @tan(value);
71 }
72
73 pub fn erf(comptime D: type, value: D.Vector) D.Vector {
74 requireFloat(D.Lane);
75 var result: D.Vector = undefined;
76 inline for (0..D.lane_count) |index| result[index] = erfScalar(D.Lane, value[index]);
77 return result;
78 }
79
80 pub fn exp(comptime D: type, value: D.Vector) D.Vector {
81 requireFloat(D.Lane);
82 return @exp(value);
83 }
84
85 pub fn exp2(comptime D: type, value: D.Vector) D.Vector {
86 requireFloat(D.Lane);
87 return @exp2(value);
88 }
89
90 pub fn expm1(comptime D: type, value: D.Vector) D.Vector {
91 return mapUnary(D, value, .expm1);
92 }
93
94 pub fn log(comptime D: type, value: D.Vector) D.Vector {
95 requireFloat(D.Lane);
96 return @log(value);
97 }
98
99 pub fn log10(comptime D: type, value: D.Vector) D.Vector {
100 return mapUnary(D, value, .log10);
101 }
102
103 pub fn log1p(comptime D: type, value: D.Vector) D.Vector {
104 return mapUnary(D, value, .log1p);
105 }
106
107 pub fn log2(comptime D: type, value: D.Vector) D.Vector {
108 requireFloat(D.Lane);
109 return @log2(value);
110 }
111
112 pub fn sin(comptime D: type, value: D.Vector) D.Vector {
113 requireFloat(D.Lane);
114 return @sin(value);
115 }
116
117 pub fn sinh(comptime D: type, value: D.Vector) D.Vector {
118 return mapUnary(D, value, .sinh);
119 }
120
121 pub fn cosh(comptime D: type, value: D.Vector) D.Vector {
122 return mapUnary(D, value, .cosh);
123 }
124
125 pub fn tanh(comptime D: type, value: D.Vector) D.Vector {
126 return mapUnary(D, value, .tanh);
127 }
128
129 pub fn tgamma(comptime D: type, value: D.Vector) D.Vector {
130 return mapUnary(D, value, .tgamma);
131 }
132
133 pub fn logGamma(comptime D: type, value: D.Vector) D.Vector {
134 return mapUnary(D, value, .log_gamma);
135 }
136
137 pub fn sinCos(
138 comptime D: type,
139 value: D.Vector,
140 sine: *D.Vector,
141 cosine: *D.Vector,
142 ) void {
143 requireFloat(D.Lane);
144 sine.* = @sin(value);
145 cosine.* = @cos(value);
146 }
147
148 pub fn hypot(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
149 return mapBinary(D, a, b, .hypot);
150 }
151
152 pub fn pow(comptime D: type, base: D.Vector, exponent: D.Vector) D.Vector {
153 return mapBinary(D, base, exponent, .pow);
154 }
155
156 fn mapUnary(comptime D: type, value: D.Vector, comptime operation: Unary) D.Vector {
157 requireFloat(D.Lane);
158 var result: D.Vector = undefined;
159 inline for (0..D.lane_count) |index| {
160 const lane = value[index];
161 result[index] = switch (operation) {
162 .acos => std.math.acos(lane),
163 .acosh => std.math.acosh(lane),
164 .asin => std.math.asin(lane),
165 .asinh => std.math.asinh(lane),
166 .atan => std.math.atan(lane),
167 .atanh => std.math.atanh(lane),
168 .cbrt => std.math.cbrt(lane),
169 .expm1 => std.math.expm1(lane),
170 .log10 => std.math.log10(lane),
171 .log1p => std.math.log1p(lane),
172 .sinh => std.math.sinh(lane),
173 .cosh => std.math.cosh(lane),
174 .tanh => std.math.tanh(lane),
175 .tgamma => std.math.gamma(D.Lane, lane),
176 .log_gamma => std.math.lgamma(D.Lane, lane),
177 };
178 }
179 return result;
180 }
181
182 fn mapBinary(
183 comptime D: type,
184 a: D.Vector,
185 b: D.Vector,
186 comptime operation: Binary,
187 ) D.Vector {
188 requireFloat(D.Lane);
189 var result: D.Vector = undefined;
190 inline for (0..D.lane_count) |index| {
191 result[index] = switch (operation) {
192 .atan2 => std.math.atan2(a[index], b[index]),
193 .hypot => std.math.hypot(a[index], b[index]),
194 .pow => std.math.pow(D.Lane, a[index], b[index]),
195 };
196 }
197 return result;
198 }
199
200 fn erfScalar(comptime T: type, input: T) T {
201 if (std.math.isNan(input)) return input;
202 const magnitude = @abs(input);
203 const limit: T = if (T == f32) 14 else 37.519379347;
204 const x = @min(magnitude, limit);
205 const z = x * x;
206 const small = if (T == f32) erfSmall32(x, z) else erfSmall64(x, z);
207 const factor = if (T == f32) erfFactor32(x) else erfFactor64(x);
208 const large = @as(T, 1) - @exp(-z) * factor;
209 return std.math.copysign(if (x < 1) small else large, input);
210 }
211
212 fn erfSmall32(x: f32, z: f32) f32 {
213 return x * polynomial(f32, z, .{
214 1.128379165726710,
215 -0.3761262582423300,
216 0.1128358514861418,
217 -0.02685381193529856,
218 0.005188327685732524,
219 -0.0008010193625184903,
220 0.00007853861353153693,
221 });
222 }
223
224 fn erfFactor32(x: f32) f32 {
225 const inverse = 1 / x;
226 const w = inverse * inverse;
227 const first = polynomial(f32, w, .{
228 0.5638259427386472,
229 -0.2741127028184656,
230 0.3404879937665872,
231 -0.4944515323274145,
232 0.6210004621745983,
233 -0.5824733027278666,
234 0.3687424674597105,
235 -0.1387039388740657,
236 0.02326819970068386,
237 });
238 const second = polynomial(f32, w, .{
239 0.5641895067754075,
240 -0.2820767439740514,
241 0.4218463358204948,
242 -1.015265279202700,
243 2.921019019210786,
244 -7.495518717768503,
245 12.97719955372516,
246 -10.47766399936249,
247 });
248 return inverse * if (x < 2) first else second;
249 }
250
251 fn erfSmall64(x: f64, z: f64) f64 {
252 const numerator = polynomial(f64, z, .{
253 55592.3013010394962768,
254 7003.32514112805075473,
255 2232.00534594684319226,
256 90.0260197203842689217,
257 9.60497373987051638749,
258 });
259 const denominator = polynomial(f64, z, .{
260 49267.3942608635921086,
261 22629.0000613890934246,
262 4594.32382970980127987,
263 521.357949780152679795,
264 33.5617141647503099647,
265 1,
266 });
267 return x * numerator / denominator;
268 }
269
270 fn erfFactor64(x: f64) f64 {
271 const first_numerator = polynomial(f64, x, .{
272 557.535335369399327526,
273 1027.55188689515710272,
274 934.528527171957607540,
275 526.445194995477358631,
276 196.520832956077098242,
277 48.6371970985681366614,
278 7.46321056442269912687,
279 0.564189564831068821977,
280 0.000000000246196981473530512524,
281 });
282 const first_denominator = polynomial(f64, x, .{
283 557.535340817727675546,
284 1656.66309194161350182,
285 2246.33760818710981792,
286 1823.90916687909736289,
287 975.708501743205489753,
288 354.937778887819891062,
289 86.7072140885989742329,
290 13.2281951154744992508,
291 1,
292 });
293 const second_numerator = polynomial(f64, x, .{
294 2.97886665372100240670,
295 7.40974269950448939160,
296 6.16021097993053585195,
297 5.01905042251180477414,
298 1.27536670759978104416,
299 0.564189583547755073984,
300 });
301 const second_denominator = polynomial(f64, x, .{
302 3.36907645100081516050,
303 9.60896809063285878198,
304 17.0814450747565897222,
305 12.0489539808096656605,
306 9.39603524938001434673,
307 2.26052863220117276590,
308 1,
309 });
310 return if (x < 8)
311 first_numerator / first_denominator
312 else
313 second_numerator / second_denominator;
314 }
315
316 fn polynomial(comptime T: type, x: T, comptime coefficients: anytype) T {
317 var result: T = coefficients[coefficients.len - 1];
318 inline for (1..coefficients.len) |offset| {
319 const index = coefficients.len - 1 - offset;
320 result = @mulAdd(T, result, x, @as(T, coefficients[index]));
321 }
322 return result;
323 }
324
325 fn requireFloat(comptime T: type) void {
326 if (T != f32 and T != f64) @compileError("Highway contributed math requires f32 or f64 lanes");
327 }
328
329 fn expectNear(comptime T: type, expected: T, actual: T, tolerance: T) !void {
330 if (std.math.isNan(expected)) return std.testing.expect(std.math.isNan(actual));
331 if (std.math.isInf(expected)) return std.testing.expectEqual(expected, actual);
332 try std.testing.expectApproxEqAbs(expected, actual, tolerance * @max(1, @abs(expected)));
333 }
334
335 test "Highway contributed elementary math matches scalar references" {
336 const simd = @import("root.zig");
337 inline for (.{ f32, f64 }) |T| {
338 const D = simd.FixedTag(T, 4);
339 const angles: D.Vector = .{ -1, -0.25, 0.25, 1 };
340 const sine = sin(D, angles);
341 const cosine = cos(D, angles);
342 const tangent = tan(D, angles);
343 inline for (0..D.lane_count) |index| {
344 try std.testing.expectApproxEqRel(@sin(angles[index]), sine[index], std.math.floatEps(T) * 2);
345 try std.testing.expectApproxEqRel(@cos(angles[index]), cosine[index], std.math.floatEps(T) * 2);
346 try std.testing.expectApproxEqRel(@tan(angles[index]), tangent[index], std.math.floatEps(T) * 2);
347 }
348
349 const positive: D.Vector = .{ 0.25, 0.5, 2, 4 };
350 const exponent = exp(D, angles);
351 const exponent2 = exp2(D, angles);
352 const natural_log = log(D, positive);
353 inline for (0..D.lane_count) |index| {
354 try std.testing.expectApproxEqRel(@exp(angles[index]), exponent[index], std.math.floatEps(T) * 2);
355 try std.testing.expectApproxEqRel(@exp2(angles[index]), exponent2[index], std.math.floatEps(T) * 2);
356 try std.testing.expectApproxEqAbs(@log(positive[index]), natural_log[index], std.math.floatEps(T) * 2);
357 }
358 }
359 }
360
361 test "Highway contributed inverse hyperbolic and paired functions preserve identities" {
362 const simd = @import("root.zig");
363 const D = simd.FixedTag(f64, 4);
364 const unit: D.Vector = .{ -0.75, -0.25, 0.25, 0.75 };
365 const asin_result = asin(D, unit);
366 const acos_result = acos(D, unit);
367 const atan_result = atan(D, unit);
368 const atanh_result = atanh(D, unit);
369 inline for (0..D.lane_count) |index| {
370 try std.testing.expectApproxEqAbs(unit[index], @sin(asin_result[index]), 0x1p-50);
371 try std.testing.expectApproxEqAbs(@as(f64, std.math.pi / 2.0), asin_result[index] + acos_result[index], 0x1p-50);
372 try std.testing.expectApproxEqAbs(unit[index], @tan(atan_result[index]), 0x1p-50);
373 try std.testing.expectApproxEqAbs(unit[index], std.math.tanh(atanh_result[index]), 0x1p-50);
374 }
375
376 var sine: D.Vector = undefined;
377 var cosine: D.Vector = undefined;
378 sinCos(D, unit, &sine, &cosine);
379 try std.testing.expect(@reduce(.And, sine == sin(D, unit)));
380 try std.testing.expect(@reduce(.And, cosine == cos(D, unit)));
381 }
382
383 test "Highway contributed roots error function gamma power and hypot cover reference values" {
384 const simd = @import("root.zig");
385 const D = simd.FixedTag(f64, 4);
386 const roots = cbrt(D, @as(D.Vector, .{ -8, -1, 1, 27 }));
387 try std.testing.expectApproxEqAbs(@as(f64, -2), roots[0], 0x1p-50);
388 try std.testing.expectApproxEqAbs(@as(f64, 3), roots[3], 0x1p-50);
389
390 const errors = erf(D, @as(D.Vector, .{ -2, -1, 0, 1 }));
391 try std.testing.expectApproxEqAbs(@as(f64, -0.9953222650189527), errors[0], 0x1p-50);
392 try std.testing.expectApproxEqAbs(@as(f64, -0.8427007929497149), errors[1], 0x1p-50);
393 try std.testing.expectEqual(@as(f64, 0), errors[2]);
394 try std.testing.expectApproxEqAbs(@as(f64, 0.8427007929497149), errors[3], 0x1p-50);
395
396 const gamma = tgamma(D, @as(D.Vector, .{ 0.5, 1, 4, 5 }));
397 try std.testing.expectApproxEqAbs(@sqrt(std.math.pi), gamma[0], 0x1p-48);
398 try std.testing.expectApproxEqAbs(@as(f64, 24), gamma[3], 0x1p-48);
399 const log_gamma = logGamma(D, @as(D.Vector, .{ 1, 2, 4, 5 }));
400 try std.testing.expectApproxEqAbs(@log(@as(f64, 24)), log_gamma[3], 0x1p-48);
401
402 const powers = pow(D, @as(D.Vector, .{ 2, 4, 9, 16 }), @as(D.Vector, .{ 3, 0.5, 0.5, -1 }));
403 try std.testing.expect(@reduce(.And, powers == @as(D.Vector, .{ 8, 2, 3, 0.0625 })));
404 const lengths = hypot(D, @as(D.Vector, .{ 3, 5, 8, 7 }), @as(D.Vector, .{ 4, 12, 15, 24 }));
405 try std.testing.expect(@reduce(.And, lengths == @as(D.Vector, .{ 5, 13, 17, 25 })));
406 }
407
408 test "Highway contributed logarithmic and hyperbolic families cover both precisions" {
409 const simd = @import("root.zig");
410 inline for (.{ f32, f64 }) |T| {
411 const D = simd.FixedTag(T, 4);
412 const values: D.Vector = .{ 0.125, 0.5, 2, 8 };
413 const signed_values: D.Vector = .{ -2, -0.5, 0.5, 2 };
414 const epsilon = std.math.floatEps(T) * 8;
415
416 const base_ten = log10(D, values);
417 const base_two = log2(D, values);
418 const one_plus = log1p(D, values);
419 const exp_minus_one = expm1(D, signed_values);
420 const hyperbolic_sine = sinh(D, signed_values);
421 const hyperbolic_cosine = cosh(D, signed_values);
422 const inverse_sine = asinh(D, signed_values);
423 const inverse_cosine = acosh(D, values + @as(D.Vector, @splat(1)));
424 inline for (0..D.lane_count) |index| {
425 try std.testing.expectApproxEqAbs(std.math.log10(values[index]), base_ten[index], epsilon);
426 try std.testing.expectApproxEqAbs(std.math.log2(values[index]), base_two[index], epsilon);
427 try std.testing.expectApproxEqAbs(std.math.log1p(values[index]), one_plus[index], epsilon);
428 try std.testing.expectApproxEqRel(std.math.expm1(signed_values[index]), exp_minus_one[index], epsilon);
429 try std.testing.expectApproxEqRel(std.math.sinh(signed_values[index]), hyperbolic_sine[index], epsilon);
430 try std.testing.expectApproxEqRel(std.math.cosh(signed_values[index]), hyperbolic_cosine[index], epsilon);
431 try std.testing.expectApproxEqAbs(signed_values[index], std.math.sinh(inverse_sine[index]), epsilon * 2);
432 try std.testing.expectApproxEqAbs(values[index] + 1, std.math.cosh(inverse_cosine[index]), epsilon * 4);
433 }
434
435 const errors = erf(D, signed_values);
436 try std.testing.expectApproxEqAbs(@as(T, -0.9953222650189527), errors[0], epsilon * 4);
437 try std.testing.expectApproxEqAbs(@as(T, 0.9953222650189527), errors[3], epsilon * 4);
438 _ = cbrtNormal(D, values);
439 }
440 }
441
442 test "Highway contributed atan2 preserves signed zero infinity and NaN" {
443 const simd = @import("root.zig");
444 const D = simd.FixedTag(f64, 4);
445 const y: D.Vector = .{ 0.0, -0.0, std.math.inf(f64), std.math.nan(f64) };
446 const x: D.Vector = .{ -1, -1, std.math.inf(f64), 1 };
447 const result = atan2(D, y, x);
448 try std.testing.expectEqual(std.math.pi, result[0]);
449 try std.testing.expectEqual(-std.math.pi, result[1]);
450 try std.testing.expectApproxEqAbs(@as(f64, std.math.pi / 4.0), result[2], 0x1p-52);
451 try std.testing.expect(std.math.isNan(result[3]));
452 }
453
454 test "Highway contributed math samples every upstream numerical family" {
455 const simd = @import("root.zig");
456 inline for (.{ f32, f64 }) |T| {
457 const D = simd.FixedTag(T, 4);
458 const tolerance: T = std.math.floatEps(T) * 16;
459 var sample: usize = 0;
460 while (sample < 256) : (sample += 1) {
461 var unit_array: [D.lane_count]T = undefined;
462 inline for (0..D.lane_count) |lane_index| {
463 const ordinal = sample * D.lane_count + lane_index;
464 const fraction = @as(T, @floatFromInt(ordinal * 2 + 1)) / @as(T, 2048);
465 unit_array[lane_index] = fraction * 2 - 1;
466 }
467 const unit: D.Vector = unit_array;
468 const positive = (unit + @as(D.Vector, @splat(1))) * @as(D.Vector, @splat(3.9375)) + @as(D.Vector, @splat(0.125));
469 const moderate = unit * @as(D.Vector, @splat(8));
470 const angle = unit * @as(D.Vector, @splat(39_000));
471 const near_unit = unit * @as(D.Vector, @splat(0.99));
472 const gamma_input = (unit + @as(D.Vector, @splat(1))) * @as(D.Vector, @splat(17.25)) + @as(D.Vector, @splat(0.5));
473 const power_exponent = unit * @as(D.Vector, @splat(2));
474
475 const outputs = .{
476 acos(D, unit),
477 asin(D, unit),
478 atan(D, angle),
479 acosh(D, positive + @as(D.Vector, @splat(1))),
480 asinh(D, moderate),
481 atanh(D, near_unit),
482 cbrt(D, angle),
483 cos(D, angle),
484 tan(D, angle),
485 exp(D, moderate),
486 exp2(D, moderate),
487 expm1(D, moderate),
488 log(D, positive),
489 log10(D, positive),
490 log1p(D, positive),
491 log2(D, positive),
492 sin(D, angle),
493 sinh(D, moderate),
494 cosh(D, moderate),
495 tanh(D, moderate),
496 tgamma(D, gamma_input),
497 logGamma(D, gamma_input),
498 hypot(D, angle, moderate),
499 pow(D, positive, power_exponent),
500 };
501
502 inline for (0..D.lane_count) |lane_index| {
503 const x = unit[lane_index];
504 const p = positive[lane_index];
505 const m = moderate[lane_index];
506 const a = angle[lane_index];
507 const g = gamma_input[lane_index];
508 const e = power_exponent[lane_index];
509 const references = .{
510 std.math.acos(x),
511 std.math.asin(x),
512 std.math.atan(a),
513 std.math.acosh(p + 1),
514 std.math.asinh(m),
515 std.math.atanh(x * 0.99),
516 std.math.cbrt(a),
517 @cos(a),
518 @tan(a),
519 @exp(m),
520 @exp2(m),
521 std.math.expm1(m),
522 @log(p),
523 std.math.log10(p),
524 std.math.log1p(p),
525 std.math.log2(p),
526 @sin(a),
527 std.math.sinh(m),
528 std.math.cosh(m),
529 std.math.tanh(m),
530 std.math.gamma(T, g),
531 std.math.lgamma(T, g),
532 std.math.hypot(a, m),
533 std.math.pow(T, p, e),
534 };
535 inline for (outputs, references) |output, expected| {
536 try expectNear(T, expected, output[lane_index], tolerance);
537 }
538 }
539 }
540 }
541 }
542
543 test "Highway contributed hypot and power preserve special values" {
544 const simd = @import("root.zig");
545 inline for (.{ f32, f64 }) |T| {
546 const D = simd.FixedTag(T, 4);
547 const infinity = std.math.inf(T);
548 const not_number = std.math.nan(T);
549 const lengths = hypot(
550 D,
551 @as(D.Vector, .{ infinity, not_number, std.math.floatMax(T), 0 }),
552 @as(D.Vector, .{ not_number, infinity, std.math.floatMax(T), -0.0 }),
553 );
554 try std.testing.expect(std.math.isInf(lengths[0]));
555 try std.testing.expect(std.math.isInf(lengths[1]));
556 try std.testing.expect(std.math.isInf(lengths[2]));
557 try std.testing.expectEqual(@as(T, 0), lengths[3]);
558
559 const powers = pow(
560 D,
561 @as(D.Vector, .{ -2, -2, -2, -0.0 }),
562 @as(D.Vector, .{ 3, 4, 0.5, -3 }),
563 );
564 try std.testing.expectEqual(@as(T, -8), powers[0]);
565 try std.testing.expectEqual(@as(T, 16), powers[1]);
566 try std.testing.expect(std.math.isNan(powers[2]));
567 try std.testing.expect(std.math.isNegativeInf(powers[3]));
568 }
569 }