lib/simd/src/shift.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn shiftLeft(
4 comptime D: type,
5 comptime amount: usize,
6 value: D.Vector,
7 ) D.Vector {
8 validate(D);
9 validateAmount(D, amount);
10 const bits: UnsignedVector(D) = @bitCast(value);
11 return @bitCast(bits << sameAmount(D, amount));
12 }
13
14 pub fn shiftRight(
15 comptime D: type,
16 comptime amount: usize,
17 value: D.Vector,
18 ) D.Vector {
19 validate(D);
20 validateAmount(D, amount);
21 return value >> sameAmount(D, amount);
22 }
23
24 pub fn shiftLeftSame(comptime D: type, value: D.Vector, amount: usize) D.Vector {
25 validate(D);
26 std.debug.assert(amount < @bitSizeOf(D.Lane));
27 const bits: UnsignedVector(D) = @bitCast(value);
28 return @bitCast(bits << sameAmount(D, amount));
29 }
30
31 pub fn shiftRightSame(comptime D: type, value: D.Vector, amount: usize) D.Vector {
32 validate(D);
33 std.debug.assert(amount < @bitSizeOf(D.Lane));
34 return value >> sameAmount(D, amount);
35 }
36
37 pub fn shl(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {
38 validate(D);
39 const bits: UnsignedVector(D) = @bitCast(value);
40 return @bitCast(bits << laneAmounts(D, amounts));
41 }
42
43 pub fn shr(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {
44 validate(D);
45 return value >> laneAmounts(D, amounts);
46 }
47
48 pub fn averageRound(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
49 validate(D);
50 return (a | b) -% ((a ^ b) >> sameAmount(D, 1));
51 }
52
53 pub fn roundingShiftRight(
54 comptime D: type,
55 comptime amount: usize,
56 value: D.Vector,
57 ) D.Vector {
58 validateAmount(D, amount);
59 if (amount == 0) return value;
60 return averageRound(D, shiftRight(D, amount - 1, value), @splat(0));
61 }
62
63 pub fn roundingShiftRightSame(
64 comptime D: type,
65 value: D.Vector,
66 amount: usize,
67 ) D.Vector {
68 std.debug.assert(amount < @bitSizeOf(D.Lane));
69 if (amount == 0) return value;
70 return averageRound(D, shiftRightSame(D, value, amount - 1), @splat(0));
71 }
72
73 pub fn roundingShr(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {
74 validate(D);
75 const U = UnsignedVector(D);
76 const raw: U = @bitCast(amounts);
77 const zero: U = @splat(0);
78 const one: U = @splat(1);
79 const scaled_amounts: U = @select(UnsignedLane(D), raw == zero, zero, raw -% one);
80 const scaled = shr(D, value, @bitCast(scaled_amounts));
81 const other = @select(D.Lane, raw == zero, scaled, @as(D.Vector, @splat(0)));
82 return averageRound(D, scaled, other);
83 }
84
85 pub fn maskedShiftLeftOr(
86 comptime D: type,
87 comptime amount: usize,
88 no: D.Vector,
89 mask: D.Mask,
90 value: D.Vector,
91 ) D.Vector {
92 return @select(D.Lane, mask, shiftLeft(D, amount, value), no);
93 }
94
95 pub fn maskedShiftLeft(
96 comptime D: type,
97 comptime amount: usize,
98 mask: D.Mask,
99 value: D.Vector,
100 ) D.Vector {
101 return maskedShiftLeftOr(D, amount, @splat(0), mask, value);
102 }
103
104 pub fn maskedShiftRightOr(
105 comptime D: type,
106 comptime amount: usize,
107 no: D.Vector,
108 mask: D.Mask,
109 value: D.Vector,
110 ) D.Vector {
111 return @select(D.Lane, mask, shiftRight(D, amount, value), no);
112 }
113
114 pub fn maskedShiftRight(
115 comptime D: type,
116 comptime amount: usize,
117 mask: D.Mask,
118 value: D.Vector,
119 ) D.Vector {
120 return maskedShiftRightOr(D, amount, @splat(0), mask, value);
121 }
122
123 pub fn maskedShlOr(
124 comptime D: type,
125 no: D.Vector,
126 mask: D.Mask,
127 value: D.Vector,
128 amounts: D.Vector,
129 ) D.Vector {
130 return @select(D.Lane, mask, shl(D, value, amounts), no);
131 }
132
133 pub fn maskedShl(
134 comptime D: type,
135 mask: D.Mask,
136 value: D.Vector,
137 amounts: D.Vector,
138 ) D.Vector {
139 return maskedShlOr(D, @splat(0), mask, value, amounts);
140 }
141
142 pub fn maskedShrOr(
143 comptime D: type,
144 no: D.Vector,
145 mask: D.Mask,
146 value: D.Vector,
147 amounts: D.Vector,
148 ) D.Vector {
149 return @select(D.Lane, mask, shr(D, value, amounts), no);
150 }
151
152 pub fn maskedShr(
153 comptime D: type,
154 mask: D.Mask,
155 value: D.Vector,
156 amounts: D.Vector,
157 ) D.Vector {
158 return maskedShrOr(D, @splat(0), mask, value, amounts);
159 }
160
161 fn UnsignedLane(comptime D: type) type {
162 return @Int(.unsigned, @bitSizeOf(D.Lane));
163 }
164
165 fn UnsignedVector(comptime D: type) type {
166 return @Vector(D.lane_count, UnsignedLane(D));
167 }
168
169 fn AmountVector(comptime D: type) type {
170 return @Vector(D.lane_count, std.math.Log2Int(UnsignedLane(D)));
171 }
172
173 fn sameAmount(comptime D: type, amount: usize) AmountVector(D) {
174 return @splat(@intCast(amount));
175 }
176
177 fn laneAmounts(comptime D: type, amounts: D.Vector) AmountVector(D) {
178 const raw: UnsignedVector(D) = @bitCast(amounts);
179 return @truncate(raw);
180 }
181
182 fn validate(comptime D: type) void {
183 if (@typeInfo(D.Lane) != .int) @compileError("shifts require integer lanes");
184 }
185
186 fn validateAmount(comptime D: type, comptime amount: usize) void {
187 validate(D);
188 if (amount >= @bitSizeOf(D.Lane)) @compileError("shift amount exceeds lane width");
189 }
190
191 fn verifyLaneType(comptime T: type) !void {
192 const simd = @import("root.zig");
193 const D = simd.FixedTag(T, 4);
194 const U = @Int(.unsigned, @bitSizeOf(T));
195 const UV = @Vector(4, U);
196 const raw: UV = .{ 0, 1, @as(U, 1) << (@bitSizeOf(T) - 1), std.math.maxInt(U) };
197 const value: D.Vector = @bitCast(raw);
198 const amounts: D.Vector = @splat(@as(T, 1));
199 try std.testing.expect(@reduce(.And, shiftLeft(D, 1, value) == shl(D, value, amounts)));
200 try std.testing.expect(@reduce(.And, shiftRight(D, 1, value) == shr(D, value, amounts)));
201 try std.testing.expect(@reduce(.And, shiftLeftSame(D, value, 1) == shiftLeft(D, 1, value)));
202 try std.testing.expect(@reduce(.And, shiftRightSame(D, value, 1) == shiftRight(D, 1, value)));
203 try std.testing.expect(@reduce(.And, roundingShiftRight(D, 0, value) == value));
204 }
205
206 test "Highway shifts instantiate every integer lane type" {
207 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64 }) |T| {
208 try verifyLaneType(T);
209 }
210 }
211
212 test "Highway fixed and variable shifts preserve signed semantics" {
213 const simd = @import("root.zig");
214 const D = simd.FixedTag(i8, 4);
215 const value: D.Vector = .{ -128, -3, 2, 127 };
216 const amounts: D.Vector = .{ 0, 1, 2, 7 };
217 try std.testing.expect(@reduce(.And, shiftLeft(D, 1, value) ==
218 @as(D.Vector, .{ 0, -6, 4, -2 })));
219 try std.testing.expect(@reduce(.And, shiftRight(D, 1, value) ==
220 @as(D.Vector, .{ -64, -2, 1, 63 })));
221 try std.testing.expect(@reduce(.And, shl(D, value, amounts) ==
222 @as(D.Vector, .{ -128, -6, 8, -128 })));
223 try std.testing.expect(@reduce(.And, shr(D, value, amounts) ==
224 @as(D.Vector, .{ -128, -2, 0, 0 })));
225 }
226
227 test "Highway rounding shifts round upward at the half bit" {
228 const simd = @import("root.zig");
229 const DU = simd.FixedTag(u8, 4);
230 const DI = simd.FixedTag(i8, 4);
231 try std.testing.expect(@reduce(.And, roundingShiftRight(DU, 1, @as(DU.Vector, .{ 0, 1, 2, 255 })) ==
232 @as(DU.Vector, .{ 0, 1, 1, 128 })));
233 try std.testing.expect(@reduce(.And, roundingShiftRight(DI, 1, @as(DI.Vector, .{ -128, -3, 2, 127 })) ==
234 @as(DI.Vector, .{ -64, -1, 1, 64 })));
235 const value: DU.Vector = .{ 0, 3, 7, 255 };
236 const amounts: DU.Vector = .{ 0, 1, 2, 7 };
237 try std.testing.expect(@reduce(.And, roundingShiftRightSame(DU, value, 2) ==
238 @as(DU.Vector, .{ 0, 1, 2, 64 })));
239 try std.testing.expect(@reduce(.And, roundingShr(DU, value, amounts) ==
240 @as(DU.Vector, .{ 0, 2, 2, 2 })));
241 }
242
243 test "Highway masked shifts preserve inactive lanes" {
244 const simd = @import("root.zig");
245 const D = simd.FixedTag(u16, 4);
246 const value: D.Vector = .{ 1, 2, 3, 4 };
247 const mask: D.Mask = .{ true, false, true, false };
248 const no: D.Vector = @splat(9);
249 try std.testing.expect(@reduce(.And, maskedShiftLeftOr(D, 2, no, mask, value) ==
250 @as(D.Vector, .{ 4, 9, 12, 9 })));
251 try std.testing.expect(@reduce(.And, maskedShiftRight(D, 1, mask, value) ==
252 @as(D.Vector, .{ 0, 0, 1, 0 })));
253 try std.testing.expect(@reduce(.And, maskedShiftLeft(D, 2, mask, value) ==
254 @as(D.Vector, .{ 4, 0, 12, 0 })));
255 const amounts: D.Vector = .{ 0, 1, 2, 3 };
256 try std.testing.expect(@reduce(.And, maskedShl(D, mask, value, amounts) ==
257 @as(D.Vector, .{ 1, 0, 12, 0 })));
258 try std.testing.expect(@reduce(.And, maskedShrOr(D, no, mask, value, amounts) ==
259 @as(D.Vector, .{ 1, 9, 0, 9 })));
260 try std.testing.expect(@reduce(.And, maskedShr(D, mask, value, amounts) ==
261 @as(D.Vector, .{ 1, 0, 0, 0 })));
262 }