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 }