lib/simd/src/floating.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn sqrt(comptime D: type, value: D.Vector) D.Vector {
4 requireFloat(D.Lane, "sqrt");
5 return @sqrt(value);
6 }
7
8 pub fn approximateReciprocalSqrt(comptime D: type, value: D.Vector) D.Vector {
9 requireFloat(D.Lane, "approximateReciprocalSqrt");
10 return @as(D.Vector, @splat(1)) / @sqrt(value);
11 }
12
13 pub fn approximateReciprocal(comptime D: type, value: D.Vector) D.Vector {
14 requireFloat(D.Lane, "approximateReciprocal");
15 return @as(D.Vector, @splat(1)) / value;
16 }
17
18 pub fn getExponent(comptime D: type, value: D.Vector) D.Vector {
19 requireFloat(D.Lane, "getExponent");
20 const U = @Int(.unsigned, @bitSizeOf(D.Lane));
21 const mantissa_bits = mantissaBits(D.Lane);
22 const exponent_mask = exponentMask(D.Lane);
23 const bias = exponentBias(D.Lane);
24 var result: D.Vector = undefined;
25 inline for (0..D.lane_count) |index| {
26 const bits: U = @bitCast(value[index]);
27 const biased = (bits >> mantissa_bits) & exponent_mask;
28 result[index] = @floatFromInt(@as(i32, @intCast(biased)) - bias);
29 }
30 return result;
31 }
32
33 pub fn getBiasedExponent(
34 comptime D: type,
35 value: D.Vector,
36 ) D.rebind(@Int(.unsigned, @bitSizeOf(D.Lane))).Vector {
37 requireFloat(D.Lane, "getBiasedExponent");
38 const U = @Int(.unsigned, @bitSizeOf(D.Lane));
39 const DU = D.rebind(U);
40 const mantissa_bits = mantissaBits(D.Lane);
41 const exponent_mask = exponentMask(D.Lane);
42 var result: DU.Vector = undefined;
43 inline for (0..D.lane_count) |index| {
44 const bits: U = @bitCast(value[index]);
45 result[index] = (bits >> mantissa_bits) & exponent_mask;
46 }
47 return result;
48 }
49
50 pub fn round(comptime D: type, value: D.Vector) D.Vector {
51 requireFloat(D.Lane, "round");
52 var result: D.Vector = undefined;
53 inline for (0..D.lane_count) |index| {
54 result[index] = roundEvenScalar(D.Lane, value[index]);
55 }
56 return result;
57 }
58
59 pub fn trunc(comptime D: type, value: D.Vector) D.Vector {
60 requireFloat(D.Lane, "trunc");
61 return @trunc(value);
62 }
63
64 pub fn ceil(comptime D: type, value: D.Vector) D.Vector {
65 requireFloat(D.Lane, "ceil");
66 return @ceil(value);
67 }
68
69 pub fn floor(comptime D: type, value: D.Vector) D.Vector {
70 requireFloat(D.Lane, "floor");
71 return @floor(value);
72 }
73
74 pub fn maskedSqrtOr(comptime D: type, no: D.Vector, mask: D.Mask, value: D.Vector) D.Vector {
75 return @select(D.Lane, mask, sqrt(D, value), no);
76 }
77
78 pub fn maskedSqrt(comptime D: type, mask: D.Mask, value: D.Vector) D.Vector {
79 return @select(D.Lane, mask, sqrt(D, value), @as(D.Vector, @splat(0)));
80 }
81
82 pub fn maskedApproximateReciprocalSqrt(comptime D: type, mask: D.Mask, value: D.Vector) D.Vector {
83 return @select(
84 D.Lane,
85 mask,
86 approximateReciprocalSqrt(D, value),
87 @as(D.Vector, @splat(0)),
88 );
89 }
90
91 pub fn maskedApproximateReciprocal(comptime D: type, mask: D.Mask, value: D.Vector) D.Vector {
92 return @select(
93 D.Lane,
94 mask,
95 approximateReciprocal(D, value),
96 @as(D.Vector, @splat(0)),
97 );
98 }
99
100 fn roundEvenScalar(comptime T: type, value: T) T {
101 const rounded = @round(value);
102 const tie = @abs(value - @trunc(value)) == @as(T, 0.5);
103 const half = rounded / @as(T, 2);
104 if (tie and @trunc(half) != half) {
105 return rounded - if (value < 0) @as(T, -1) else @as(T, 1);
106 }
107 return rounded;
108 }
109
110 fn mantissaBits(comptime T: type) comptime_int {
111 return switch (T) {
112 f16 => 10,
113 f32 => 23,
114 f64 => 52,
115 else => @compileError("unsupported floating-point lane"),
116 };
117 }
118
119 fn exponentMask(comptime T: type) @Int(.unsigned, @bitSizeOf(T)) {
120 return switch (T) {
121 f16 => 0x1f,
122 f32 => 0xff,
123 f64 => 0x7ff,
124 else => @compileError("unsupported floating-point lane"),
125 };
126 }
127
128 fn exponentBias(comptime T: type) i32 {
129 return switch (T) {
130 f16 => 15,
131 f32 => 127,
132 f64 => 1023,
133 else => @compileError("unsupported floating-point lane"),
134 };
135 }
136
137 fn requireFloat(comptime T: type, comptime operation: []const u8) void {
138 if (comptime @typeInfo(T) != .float) @compileError(operation ++ " requires floating-point lanes");
139 }
140
141 test "Highway floating roots reciprocals and masks preserve active lanes" {
142 const simd = @import("root.zig");
143 const D = simd.FixedTag(f32, 4);
144 const value: D.Vector = .{ 1, 4, 16, 64 };
145 try std.testing.expect(@reduce(.And, sqrt(D, value) == @as(D.Vector, .{ 1, 2, 4, 8 })));
146 try std.testing.expect(@reduce(.And, approximateReciprocal(D, value) ==
147 @as(D.Vector, .{ 1, 0.25, 0.0625, 0.015625 })));
148 const reciprocal_root = approximateReciprocalSqrt(D, value);
149 try std.testing.expectApproxEqAbs(@as(f32, 0.5), reciprocal_root[1], 0.000001);
150 const mask: D.Mask = .{ true, false, true, false };
151 try std.testing.expect(@reduce(.And, maskedSqrtOr(D, @splat(9), mask, value) ==
152 @as(D.Vector, .{ 1, 9, 4, 9 })));
153 try std.testing.expect(@reduce(.And, maskedApproximateReciprocal(D, mask, value) ==
154 @as(D.Vector, .{ 1, 0, 0.0625, 0 })));
155 }
156
157 test "Highway floating rounding uses ties-to-even and preserves infinities" {
158 const simd = @import("root.zig");
159 inline for (.{ f16, f32, f64 }) |T| {
160 const D = simd.FixedTag(T, 8);
161 const value: D.Vector = .{ -2.5, -1.5, -0.4, 0.4, 1.5, 2.5, std.math.inf(T), -std.math.inf(T) };
162 const rounded = round(D, value);
163 try std.testing.expectEqual(@as(T, -2), rounded[0]);
164 try std.testing.expectEqual(@as(T, -2), rounded[1]);
165 try std.testing.expectEqual(@as(T, 0), rounded[2]);
166 try std.testing.expectEqual(@as(T, 0), rounded[3]);
167 try std.testing.expectEqual(@as(T, 2), rounded[4]);
168 try std.testing.expectEqual(@as(T, 2), rounded[5]);
169 try std.testing.expect(std.math.isInf(rounded[6]));
170 try std.testing.expect(std.math.isInf(rounded[7]));
171 _ = trunc(D, value);
172 _ = ceil(D, value);
173 _ = floor(D, value);
174 }
175 }
176
177 test "Highway exponent extraction returns unbiased floats and raw fields" {
178 const simd = @import("root.zig");
179 inline for (.{ f16, f32, f64 }) |T| {
180 const D = simd.FixedTag(T, 4);
181 const value: D.Vector = .{ 1, 2, 3, 8 };
182 try std.testing.expect(@reduce(.And, getExponent(D, value) ==
183 @as(D.Vector, .{ 0, 1, 1, 3 })));
184 const biased = getBiasedExponent(D, value);
185 const U = @Int(.unsigned, @bitSizeOf(T));
186 const bias: U = @intCast(exponentBias(T));
187 try std.testing.expectEqual(bias, biased[0]);
188 try std.testing.expectEqual(bias + 1, biased[1]);
189 try std.testing.expectEqual(bias + 3, biased[3]);
190 }
191 }