lib/simd/src/complex.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn complexConj(comptime D: type, value: D.Vector) D.Vector {
4 validate(D);
5 var result = value;
6 inline for (0..D.lane_count / 2) |index| {
7 result[index * 2 + 1] = -value[index * 2 + 1];
8 }
9 return result;
10 }
11
12 pub fn mulComplex(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
13 return multiply(D, a, b, false, null);
14 }
15
16 pub fn mulComplexConj(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
17 return multiply(D, a, b, true, null);
18 }
19
20 pub fn mulComplexAdd(comptime D: type, a: D.Vector, b: D.Vector, c: D.Vector) D.Vector {
21 return multiply(D, a, b, false, c);
22 }
23
24 pub fn mulComplexConjAdd(comptime D: type, a: D.Vector, b: D.Vector, c: D.Vector) D.Vector {
25 return multiply(D, a, b, true, c);
26 }
27
28 pub fn maskedMulComplexConjAdd(
29 comptime D: type,
30 mask: D.Mask,
31 a: D.Vector,
32 b: D.Vector,
33 c: D.Vector,
34 ) D.Vector {
35 return @select(
36 D.Lane,
37 mask,
38 mulComplexConjAdd(D, a, b, c),
39 @as(D.Vector, @splat(0)),
40 );
41 }
42
43 pub fn maskedMulComplexConj(
44 comptime D: type,
45 mask: D.Mask,
46 a: D.Vector,
47 b: D.Vector,
48 ) D.Vector {
49 return @select(
50 D.Lane,
51 mask,
52 mulComplexConj(D, a, b),
53 @as(D.Vector, @splat(0)),
54 );
55 }
56
57 pub fn maskedMulComplexOr(
58 comptime D: type,
59 no: D.Vector,
60 mask: D.Mask,
61 a: D.Vector,
62 b: D.Vector,
63 ) D.Vector {
64 return @select(D.Lane, mask, mulComplex(D, a, b), no);
65 }
66
67 fn multiply(
68 comptime D: type,
69 a: D.Vector,
70 b: D.Vector,
71 comptime conjugate_b: bool,
72 addend: ?D.Vector,
73 ) D.Vector {
74 validate(D);
75 var result: D.Vector = undefined;
76 inline for (0..D.lane_count / 2) |index| {
77 const real = index * 2;
78 const imaginary = real + 1;
79 const c_real = if (addend) |c| c[real] else 0;
80 const c_imaginary = if (addend) |c| c[imaginary] else 0;
81 if (conjugate_b) {
82 result[real] = @mulAdd(D.Lane, a[imaginary], b[imaginary], @mulAdd(D.Lane, a[real], b[real], c_real));
83 result[imaginary] = @mulAdd(D.Lane, a[imaginary], b[real], @mulAdd(D.Lane, -a[real], b[imaginary], c_imaginary));
84 } else {
85 result[real] = @mulAdd(D.Lane, -a[imaginary], b[imaginary], @mulAdd(D.Lane, a[real], b[real], c_real));
86 result[imaginary] = @mulAdd(D.Lane, a[imaginary], b[real], @mulAdd(D.Lane, a[real], b[imaginary], c_imaginary));
87 }
88 }
89 return result;
90 }
91
92 fn validate(comptime D: type) void {
93 if (comptime @typeInfo(D.Lane) != .float or D.lane_count & 1 != 0) {
94 @compileError("complex arithmetic requires an even number of floating-point lanes");
95 }
96 }
97
98 test "Highway complex multiply conjugate add and masks preserve lane pairs" {
99 const simd = @import("root.zig");
100 const D = simd.FixedTag(f32, 8);
101 const a: D.Vector = .{ 1, 2, 3, 4, -2, 5, 7, -3 };
102 const b: D.Vector = .{ 5, 6, 7, 8, 4, -1, -2, 6 };
103 try std.testing.expect(@reduce(.And, complexConj(D, a) ==
104 @as(D.Vector, .{ 1, -2, 3, -4, -2, -5, 7, 3 })));
105 try std.testing.expect(@reduce(.And, mulComplex(D, a, b) ==
106 @as(D.Vector, .{ -7, 16, -11, 52, -3, 22, 4, 48 })));
107 try std.testing.expect(@reduce(.And, mulComplexConj(D, a, b) ==
108 @as(D.Vector, .{ 17, 4, 53, 4, -13, 18, -32, -36 })));
109 try std.testing.expect(@reduce(.And, mulComplexAdd(D, a, b, @as(D.Vector, @splat(1))) ==
110 @as(D.Vector, .{ -6, 17, -10, 53, -2, 23, 5, 49 })));
111 const mask: D.Mask = .{ true, true, false, false, true, false, false, true };
112 try std.testing.expect(@reduce(.And, maskedMulComplexOr(D, @splat(9), mask, a, b) ==
113 @as(D.Vector, .{ -7, 16, 9, 9, -3, 9, 9, 48 })));
114 _ = maskedMulComplexConj(D, mask, a, b);
115 _ = maskedMulComplexConjAdd(D, mask, a, b, @splat(1));
116 }
117
118 test "Highway complex arithmetic instantiates every floating lane type" {
119 const simd = @import("root.zig");
120 inline for (.{ f16, f32, f64 }) |T| {
121 const D = simd.FixedTag(T, 4);
122 const value: D.Vector = @splat(1);
123 _ = mulComplex(D, value, value);
124 _ = mulComplexConj(D, value, value);
125 _ = mulComplexAdd(D, value, value, value);
126 _ = mulComplexConjAdd(D, value, value, value);
127 }
128 }