lib/simd/src/conditional.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const arithmetic = @import("arithmetic.zig");
3
4 pub fn ifThenElse(comptime D: type, mask: D.Mask, yes: D.Vector, no: D.Vector) D.Vector {
5 return @select(D.Lane, mask, yes, no);
6 }
7
8 pub fn ifThenElseZero(comptime D: type, mask: D.Mask, yes: D.Vector) D.Vector {
9 return ifThenElse(D, mask, yes, @as(D.Vector, @splat(0)));
10 }
11
12 pub fn ifThenZeroElse(comptime D: type, mask: D.Mask, no: D.Vector) D.Vector {
13 return ifThenElse(D, mask, @as(D.Vector, @splat(0)), no);
14 }
15
16 pub fn ifVecThenElse(comptime D: type, mask: D.Vector, yes: D.Vector, no: D.Vector) D.Vector {
17 return bitwiseIfThenElse(D, mask, yes, no);
18 }
19
20 pub fn bitwiseIfThenElse(comptime D: type, mask: D.Vector, yes: D.Vector, no: D.Vector) D.Vector {
21 const U = @Int(.unsigned, @bitSizeOf(D.Lane));
22 const DU = D.rebind(U);
23 const mask_bits: DU.Vector = @bitCast(mask);
24 const yes_bits: DU.Vector = @bitCast(yes);
25 const no_bits: DU.Vector = @bitCast(no);
26 return @bitCast((mask_bits & yes_bits) | (~mask_bits & no_bits));
27 }
28
29 pub fn zeroIfNegative(comptime D: type, value: D.Vector) D.Vector {
30 return ifThenZeroElse(D, negativeMask(D, value), value);
31 }
32
33 pub fn ifNegativeThenElse(
34 comptime D: type,
35 sign: D.Vector,
36 negative: D.Vector,
37 nonnegative: D.Vector,
38 ) D.Vector {
39 return ifThenElse(D, negativeMask(D, sign), negative, nonnegative);
40 }
41
42 pub fn ifNegativeThenElseZero(comptime D: type, sign: D.Vector, negative: D.Vector) D.Vector {
43 return ifThenElseZero(D, negativeMask(D, sign), negative);
44 }
45
46 pub fn ifNegativeThenZeroElse(comptime D: type, sign: D.Vector, nonnegative: D.Vector) D.Vector {
47 return ifThenZeroElse(D, negativeMask(D, sign), nonnegative);
48 }
49
50 pub fn ifNegativeThenNegOrUndefIfZero(
51 comptime D: type,
52 sign: D.Vector,
53 value: D.Vector,
54 ) D.Vector {
55 return ifThenElse(D, negativeMask(D, sign), arithmetic.neg(D, value), value);
56 }
57
58 fn negativeMask(comptime D: type, value: D.Vector) D.Mask {
59 return switch (@typeInfo(D.Lane)) {
60 .int => |info| if (info.signedness == .signed)
61 value < @as(D.Vector, @splat(0))
62 else
63 @compileError("negative selection requires signed integer or floating-point lanes"),
64 .float => blk: {
65 const U = @Int(.unsigned, @bitSizeOf(D.Lane));
66 const DU = D.rebind(U);
67 const bits: DU.Vector = @bitCast(value);
68 break :blk bits & @as(DU.Vector, @splat(@as(U, 1) << (@bitSizeOf(U) - 1))) !=
69 @as(DU.Vector, @splat(0));
70 },
71 else => unreachable,
72 };
73 }
74
75 test "Highway conditional selection covers masks vectors and raw bit patterns" {
76 const simd = @import("root.zig");
77 const D = simd.FixedTag(u32, 4);
78 const yes: D.Vector = .{ 1, 2, 3, 4 };
79 const no: D.Vector = .{ 10, 20, 30, 40 };
80 const mask: D.Mask = .{ true, false, true, false };
81 try std.testing.expect(@reduce(.And, ifThenElse(D, mask, yes, no) ==
82 @as(D.Vector, .{ 1, 20, 3, 40 })));
83 try std.testing.expect(@reduce(.And, ifThenElseZero(D, mask, yes) ==
84 @as(D.Vector, .{ 1, 0, 3, 0 })));
85 try std.testing.expect(@reduce(.And, ifThenZeroElse(D, mask, no) ==
86 @as(D.Vector, .{ 0, 20, 0, 40 })));
87 const vector_mask: D.Vector = .{ std.math.maxInt(u32), 0, std.math.maxInt(u32), 0 };
88 try std.testing.expect(@reduce(.And, ifVecThenElse(D, vector_mask, yes, no) ==
89 @as(D.Vector, .{ 1, 20, 3, 40 })));
90 const partial: D.Vector = @splat(0x0f0f_0f0f);
91 try std.testing.expect(@reduce(.And, bitwiseIfThenElse(D, partial, @as(D.Vector, @splat(0xaaaa_aaaa)), @as(D.Vector, @splat(0x5555_5555))) ==
92 @as(D.Vector, @splat(0x5a5a_5a5a))));
93 }
94
95 test "Highway sign conditionals use representation sign including negative zero" {
96 const simd = @import("root.zig");
97 const D = simd.FixedTag(f32, 4);
98 const sign: D.Vector = .{ -1, -0.0, 0, 1 };
99 const value: D.Vector = .{ 1, 2, 3, 4 };
100 try std.testing.expect(@reduce(.And, zeroIfNegative(D, sign) ==
101 @as(D.Vector, .{ 0, 0, 0, 1 })));
102 try std.testing.expect(@reduce(.And, ifNegativeThenElse(D, sign, value, @as(D.Vector, @splat(9))) ==
103 @as(D.Vector, .{ 1, 2, 9, 9 })));
104 try std.testing.expect(@reduce(.And, ifNegativeThenElseZero(D, sign, value) ==
105 @as(D.Vector, .{ 1, 2, 0, 0 })));
106 try std.testing.expect(@reduce(.And, ifNegativeThenZeroElse(D, sign, value) ==
107 @as(D.Vector, .{ 0, 0, 3, 4 })));
108 try std.testing.expect(@reduce(.And, ifNegativeThenNegOrUndefIfZero(D, sign, value) ==
109 @as(D.Vector, .{ -1, -2, 3, 4 })));
110 }
111
112 test "Highway sign conditionals instantiate signed integer and floating lanes" {
113 const simd = @import("root.zig");
114 inline for (.{ i8, i16, i32, i64, f16, f32, f64 }) |T| {
115 const D = simd.FixedTag(T, 4);
116 const value: D.Vector = @splat(1);
117 _ = zeroIfNegative(D, value);
118 _ = ifNegativeThenElse(D, value, value, value);
119 _ = ifNegativeThenNegOrUndefIfZero(D, value, value);
120 }
121 }