authorgravatar for 39187961+tecanec@users.noreply.github.comtecanec <39187961+tecanec@users.noreply.github.com> 2022-03-27 10:28:44+02:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2022-03-27 11:28:44+03:00
log3beef3945c87c589feebf50ed985bb6e65887110
tree59ad4dfd11f773ccff82c48be1e2a30f04ee2f3c
parenta3030221c3bb3f3d59cedc6db83c679470fca417
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

std: SIMD utility functions

This file contains a collections of functions that may be useful for SIMD, such as generating a vector with a linear range of numbers starting at zero, joining two vectors together, getting the index of the first true in a vector of bools, etc.

2 files changed, 410 insertions(+), 0 deletions(-)

lib/std/simd.zig created+409
...@@ -0,0 +1,409 @@
1//! This module provides functions for working conveniently with SIMD (Single Instruction; Multiple Data),
2//! which may offer a potential boost in performance on some targets by performing the same operations on
3//! multiple elements at once.
4//! Please be aware that some functions are known to not work on MIPS.
5
6const std = @import("std");
7const builtin = @import("builtin");
8
9pub const Vector = std.meta.Vector;
10
11pub fn suggestVectorSizeForCpu(comptime T: type, cpu: std.Target.Cpu) ?usize {
12 switch (cpu.arch) {
13 .x86_64 => {
14 // Note: This is mostly just guesswork. It'd be great if someone more qualified were to take a
15 // proper look at this.
16
17 if (T == bool and std.Target.x86.featureSetHas(.prefer_mask_registers)) return 64;
18
19 const vector_bit_size = blk: {
20 if (std.Target.x86.featureSetHas(.avx512f)) break :blk 512;
21 if (std.Target.x86.featureSetHas(.prefer_256_bit)) break :blk 256;
22 if (std.Target.x86.featureSetHas(.prefer_128_bit)) break :blk 128;
23 return null;
24 };
25 const element_bit_size = std.math.max(8, std.math.ceilPowerOfTwo(T, @bitSizeOf(T)));
26 return @divExact(vector_bit_size, element_bit_size);
27 },
28 else => @compileError("No vector sizes for this CPU architecture have yet been recommended"),
29 }
30}
31
32/// Suggests a target-dependant vector size for a given type, or null if scalars are recommended.
33/// Not yet implemented for every CPU architecture.
34pub fn suggestVectorSize(comptime T: type) ?usize {
35 return suggestVectorSizeForCpu(T, builtin.cpu);
36}
37
38fn vectorLength(comptime VectorType: type) comptime_int {
39 return switch (@typeInfo(VectorType)) {
40 .Vector => |info| info.len,
41 .Array => |info| info.len,
42 else => @compileError("Invalid type " ++ @typeName(VectorType)),
43 };
44}
45
46/// Returns the smallest type of unsigned ints capable of indexing any element within the given vector type.
47pub fn VectorIndex(comptime VectorType: type) type {
48 return std.math.IntFittingRange(0, vectorLength(VectorType) - 1);
49}
50
51/// Returns the smallest type of unsigned ints capable of holding the length of the given vector type.
52pub fn VectorCount(comptime VectorType: type) type {
53 return std.math.IntFittingRange(0, vectorLength(VectorType));
54}
55
56/// Returns a vector containing the first `len` integers in order from 0 to `len`-1.
57/// For example, `iota(i32, 8)` will return a vector containing `.{0, 1, 2, 3, 4, 5, 6, 7}`.
58pub fn iota(comptime T: type, comptime len: usize) Vector(len, T) {
59 var out: [len]T = undefined;
60 for (out) |*element, i| {
61 element.* = switch (@typeInfo(T)) {
62 .Int => @intCast(T, i),
63 .Float => @intToFloat(T, i),
64 else => @compileError("Can't use type " ++ @typeName(T) ++ " in iota."),
65 };
66 }
67 return @as(Vector(len, T), out);
68}
69
70/// Returns a vector containing the same elements as the input, but repeated until the desired length is reached.
71/// For example, `repeat(8, [_]u32{1, 2, 3})` will return a vector containing `.{1, 2, 3, 1, 2, 3, 1, 2}`.
72pub fn repeat(comptime len: usize, vec: anytype) Vector(len, std.meta.Child(@TypeOf(vec))) {
73 const Child = std.meta.Child(@TypeOf(vec));
74
75 return @shuffle(Child, vec, undefined, iota(i32, len) % @splat(len, @intCast(i32, vectorLength(@TypeOf(vec)))));
76}
77
78/// Returns a vector containing all elements of the first vector at the lower indices followed by all elements of the second vector
79/// at the higher indices.
80pub fn join(a: anytype, b: anytype) Vector(vectorLength(@TypeOf(a)) + vectorLength(@TypeOf(b)), std.meta.Child(@TypeOf(a))) {
81 const Child = std.meta.Child(@TypeOf(a));
82 const a_len = vectorLength(@TypeOf(a));
83 const b_len = vectorLength(@TypeOf(b));
84
85 return @shuffle(Child, a, b, @as([a_len]i32, iota(i32, a_len)) ++ @as([b_len]i32, ~iota(i32, b_len)));
86}
87
88/// Returns a vector whose elements alternates between those of each input vector.
89/// For example, `interlace(.{[4]u32{11, 12, 13, 14}, [4]u32{21, 22, 23, 24}})` returns a vector containing `.{11, 21, 12, 22, 13, 23, 14, 24}`.
90pub fn interlace(vecs: anytype) Vector(vectorLength(@TypeOf(vecs[0])) * vecs.len, std.meta.Child(@TypeOf(vecs[0]))) {
91 // interlace doesn't work on MIPS, for some reason.
92 // Notes from earlier debug attempt:
93 // The indices are correct. The problem seems to be with the @shuffle builtin.
94 // On MIPS, the test that interlaces small_base gives { 0, 2, 0, 0, 64, 255, 248, 200, 0, 0 }.
95 // Calling this with two inputs seems to work fine, but I'll let the compile error trigger for all inputs, just to be safe.
96 comptime if (builtin.cpu.arch.isMIPS()) @compileError("TODO: Find out why interlace() doesn't work on MIPS");
97
98 const VecType = @TypeOf(vecs[0]);
99 const vecs_arr = @as([vecs.len]VecType, vecs);
100 const Child = std.meta.Child(@TypeOf(vecs_arr[0]));
101
102 if (vecs_arr.len == 1) return vecs_arr[0];
103
104 const a_vec_count = (1 + vecs_arr.len) >> 1;
105 const b_vec_count = vecs_arr.len >> 1;
106
107 const a = interlace(@ptrCast(*const [a_vec_count]VecType, vecs_arr[0..a_vec_count]).*);
108 const b = interlace(@ptrCast(*const [b_vec_count]VecType, vecs_arr[a_vec_count..]).*);
109
110 const a_len = vectorLength(@TypeOf(a));
111 const b_len = vectorLength(@TypeOf(b));
112 const len = a_len + b_len;
113
114 const indices = comptime blk: {
115 const count_up = iota(i32, len);
116 const cycle = @divFloor(count_up, @splat(len, @intCast(i32, vecs_arr.len)));
117 const select_mask = repeat(len, join(@splat(a_vec_count, true), @splat(b_vec_count, false)));
118 const a_indices = count_up - cycle * @splat(len, @intCast(i32, b_vec_count));
119 const b_indices = shiftElementsRight(count_up - cycle * @splat(len, @intCast(i32, a_vec_count)), a_vec_count, 0);
120 break :blk @select(i32, select_mask, a_indices, ~b_indices);
121 };
122
123 return @shuffle(Child, a, b, indices);
124}
125
126/// The contents of `interlaced` is evenly split between vec_count vectors that are returned as an array. They "take turns",
127/// recieving one element from `interlaced` at a time.
128pub fn deinterlace(
129 comptime vec_count: usize,
130 interlaced: anytype,
131) [vec_count]Vector(
132 vectorLength(@TypeOf(interlaced)) / vec_count,
133 std.meta.Child(@TypeOf(interlaced)),
134) {
135 const vec_len = vectorLength(@TypeOf(interlaced)) / vec_count;
136 const Child = std.meta.Child(@TypeOf(interlaced));
137
138 var out: [vec_count]Vector(vec_len, Child) = undefined;
139
140 comptime var i: usize = 0; // for-loops don't work for this, apparently.
141 inline while (i < out.len) : (i += 1) {
142 const indices = comptime iota(i32, vec_len) * @splat(vec_len, @intCast(i32, vec_count)) + @splat(vec_len, @intCast(i32, i));
143 out[i] = @shuffle(Child, interlaced, undefined, indices);
144 }
145
146 return out;
147}
148
149pub fn extract(
150 vec: anytype,
151 comptime first: VectorIndex(@TypeOf(vec)),
152 comptime count: VectorCount(@TypeOf(vec)),
153) Vector(count, std.meta.Child(@TypeOf(vec))) {
154 const Child = std.meta.Child(@TypeOf(vec));
155 const len = vectorLength(@TypeOf(vec));
156
157 std.debug.assert(@intCast(comptime_int, first) + @intCast(comptime_int, count) <= len);
158
159 return @shuffle(Child, vec, undefined, iota(i32, count) + @splat(count, @intCast(i32, first)));
160}
161
162test "vector patterns" {
163 const base = Vector(4, u32){ 10, 20, 30, 40 };
164 const other_base = Vector(4, u32){ 55, 66, 77, 88 };
165
166 const small_bases = [5]Vector(2, u8){
167 Vector(2, u8){ 0, 1 },
168 Vector(2, u8){ 2, 3 },
169 Vector(2, u8){ 4, 5 },
170 Vector(2, u8){ 6, 7 },
171 Vector(2, u8){ 8, 9 },
172 };
173
174 try std.testing.expectEqual([6]u32{ 10, 20, 30, 40, 10, 20 }, repeat(6, base));
175 try std.testing.expectEqual([8]u32{ 10, 20, 30, 40, 55, 66, 77, 88 }, join(base, other_base));
176 try std.testing.expectEqual([2]u32{ 20, 30 }, extract(base, 1, 2));
177
178 if (comptime !builtin.cpu.arch.isMIPS()) {
179 try std.testing.expectEqual([8]u32{ 10, 55, 20, 66, 30, 77, 40, 88 }, interlace(.{ base, other_base }));
180
181 const small_braid = interlace(small_bases);
182 try std.testing.expectEqual([10]u8{ 0, 2, 4, 6, 8, 1, 3, 5, 7, 9 }, small_braid);
183 try std.testing.expectEqual(small_bases, deinterlace(small_bases.len, small_braid));
184 }
185}
186
187/// Joins two vectors, shifts them leftwards (towards lower indices) and extracts the leftmost elements into a vector the size of a and b.
188pub fn mergeShift(a: anytype, b: anytype, comptime shift: VectorCount(@TypeOf(a, b))) @TypeOf(a, b) {
189 const len = vectorLength(@TypeOf(a, b));
190
191 return extract(join(a, b), shift, len);
192}
193
194/// Elements are shifted rightwards (towards higher indices). New elements are added to the left, and the rightmost elements are cut off
195/// so that the size of the vector stays the same.
196pub fn shiftElementsRight(vec: anytype, comptime amount: VectorCount(@TypeOf(vec)), shift_in: std.meta.Child(@TypeOf(vec))) @TypeOf(vec) {
197 // It may be possible to implement shifts and rotates with a runtime-friendly slice of two joined vectors, as the length of the
198 // slice would be comptime-known. This would permit vector shifts and rotates by a non-comptime-known amount.
199 // However, I am unsure whether compiler optimizations would handle that well enough on all platforms.
200 const len = vectorLength(@TypeOf(vec));
201
202 return mergeShift(@splat(len, shift_in), vec, len - amount);
203}
204
205/// Elements are shifted leftwards (towards lower indices). New elements are added to the right, and the leftmost elements are cut off
206/// so that no elements with indices below 0 remain.
207pub fn shiftElementsLeft(vec: anytype, comptime amount: VectorCount(@TypeOf(vec)), shift_in: std.meta.Child(@TypeOf(vec))) @TypeOf(vec) {
208 const len = vectorLength(@TypeOf(vec));
209
210 return mergeShift(vec, @splat(len, shift_in), amount);
211}
212
213/// Elements are shifted leftwards (towards lower indices). Elements that leave to the left will reappear to the right in the same order.
214pub fn rotateElementsLeft(vec: anytype, comptime amount: VectorCount(@TypeOf(vec))) @TypeOf(vec) {
215 return mergeShift(vec, vec, amount);
216}
217
218/// Elements are shifted rightwards (towards higher indices). Elements that leave to the right will reappear to the left in the same order.
219pub fn rotateElementsRight(vec: anytype, comptime amount: VectorCount(@TypeOf(vec))) @TypeOf(vec) {
220 return rotateElementsLeft(vec, vectorLength(@TypeOf(vec)) - amount);
221}
222
223pub fn reverseOrder(vec: anytype) @TypeOf(vec) {
224 const Child = std.meta.Child(@TypeOf(vec));
225 const len = vectorLength(@TypeOf(vec));
226
227 return @shuffle(Child, vec, undefined, @splat(len, @intCast(i32, len) - 1) - iota(i32, len));
228}
229
230test "vector shifting" {
231 const base = Vector(4, u32){ 10, 20, 30, 40 };
232
233 try std.testing.expectEqual([4]u32{ 30, 40, 999, 999 }, shiftElementsLeft(base, 2, 999));
234 try std.testing.expectEqual([4]u32{ 999, 999, 10, 20 }, shiftElementsRight(base, 2, 999));
235 try std.testing.expectEqual([4]u32{ 20, 30, 40, 10 }, rotateElementsLeft(base, 1));
236 try std.testing.expectEqual([4]u32{ 40, 10, 20, 30 }, rotateElementsRight(base, 1));
237 try std.testing.expectEqual([4]u32{ 40, 30, 20, 10 }, reverseOrder(base));
238}
239
240pub fn firstTrue(vec: anytype) ?VectorIndex(@TypeOf(vec)) {
241 const len = vectorLength(@TypeOf(vec));
242 const IndexInt = VectorIndex(@TypeOf(vec));
243
244 if (!@reduce(.Or, vec)) {
245 return null;
246 }
247 const indices = @select(IndexInt, vec, iota(IndexInt, len), @splat(len, ~@as(IndexInt, 0)));
248 return @reduce(.Min, indices);
249}
250
251pub fn lastTrue(vec: anytype) ?VectorIndex(@TypeOf(vec)) {
252 const len = vectorLength(@TypeOf(vec));
253 const IndexInt = VectorIndex(@TypeOf(vec));
254
255 if (!@reduce(.Or, vec)) {
256 return null;
257 }
258 const indices = @select(IndexInt, vec, iota(IndexInt, len), @splat(len, @as(IndexInt, 0)));
259 return @reduce(.Max, indices);
260}
261
262pub fn countTrues(vec: anytype) VectorCount(@TypeOf(vec)) {
263 const len = vectorLength(@TypeOf(vec));
264 const CountIntType = VectorCount(@TypeOf(vec));
265
266 const one_if_true = @select(CountIntType, vec, @splat(len, @as(CountIntType, 1)), @splat(len, @as(CountIntType, 0)));
267 return @reduce(.Add, one_if_true);
268}
269
270pub fn firstIndexOfValue(vec: anytype, value: std.meta.Child(@TypeOf(vec))) ?VectorIndex(@TypeOf(vec)) {
271 const len = vectorLength(@TypeOf(vec));
272
273 return firstTrue(vec == @splat(len, value));
274}
275
276pub fn lastIndexOfValue(vec: anytype, value: std.meta.Child(@TypeOf(vec))) ?VectorIndex(@TypeOf(vec)) {
277 const len = vectorLength(@TypeOf(vec));
278
279 return lastTrue(vec == @splat(len, value));
280}
281
282pub fn countElementsWithValue(vec: anytype, value: std.meta.Child(@TypeOf(vec))) VectorCount(@TypeOf(vec)) {
283 const len = vectorLength(@TypeOf(vec));
284
285 return countTrues(vec == @splat(len, value));
286}
287
288test "vector searching" {
289 const base = Vector(8, u32){ 6, 4, 7, 4, 4, 2, 3, 7 };
290
291 try std.testing.expectEqual(@as(?u3, 1), firstIndexOfValue(base, 4));
292 try std.testing.expectEqual(@as(?u3, 4), lastIndexOfValue(base, 4));
293 try std.testing.expectEqual(@as(?u3, null), lastIndexOfValue(base, 99));
294 try std.testing.expectEqual(@as(u4, 3), countElementsWithValue(base, 4));
295}
296
297/// Same as prefixScan, but with a user-provided, mathematically associative function.
298pub fn prefixScanWithFunc(
299 comptime hop: isize,
300 vec: anytype,
301 /// The error type that `func` might return. Set this to `void` if `func` doesn't return an error union.
302 comptime ErrorType: type,
303 comptime func: fn (@TypeOf(vec), @TypeOf(vec)) if (ErrorType == void) @TypeOf(vec) else ErrorType!@TypeOf(vec),
304 /// When one operand of the operation performed by `func` is this value, the result must equal the other operand.
305 /// For example, this should be 0 for addition or 1 for multiplication.
306 comptime identity: std.meta.Child(@TypeOf(vec)),
307) if (ErrorType == void) @TypeOf(vec) else ErrorType!@TypeOf(vec) {
308 // I haven't debugged this, but it might be a cousin of sorts to what's going on with interlace.
309 comptime if (builtin.cpu.arch.isMIPS()) @compileError("TODO: Find out why prefixScan doesn't work on MIPS");
310
311 const len = vectorLength(@TypeOf(vec));
312
313 if (hop == 0) @compileError("hop can not be 0; you'd be going nowhere forever!");
314 const abs_hop = if (hop < 0) -hop else hop;
315
316 var acc = vec;
317 comptime var i = 0;
318 inline while ((abs_hop << i) < len) : (i += 1) {
319 const shifted = if (hop < 0) shiftElementsLeft(acc, abs_hop << i, identity) else shiftElementsRight(acc, abs_hop << i, identity);
320
321 acc = if (ErrorType == void) func(acc, shifted) else try func(acc, shifted);
322 }
323 return acc;
324}
325
326/// Returns a vector whose elements are the result of performing the specified operation on the corresponding
327/// element of the input vector and every hop'th element that came before it (or after, if hop is negative).
328/// Supports the same operations as the @reduce() builtin. Takes O(logN) to compute.
329/// The scan is not linear, which may affect floating point errors. This may affect the determinism of
330/// algorithms that use this function.
331pub fn prefixScan(comptime op: std.builtin.ReduceOp, comptime hop: isize, vec: anytype) @TypeOf(vec) {
332 const VecType = @TypeOf(vec);
333 const Child = std.meta.Child(VecType);
334 const len = vectorLength(VecType);
335
336 const identity = comptime switch (@typeInfo(Child)) {
337 .Bool => switch (op) {
338 .Or, .Xor => false,
339 .And => true,
340 else => @compileError("Invalid prefixScan operation " ++ @tagName(op) ++ " for vector of booleans."),
341 },
342 .Int => switch (op) {
343 .Max => std.math.minInt(Child),
344 .Add, .Or, .Xor => 0,
345 .Mul => 1,
346 .And, .Min => std.math.maxInt(Child),
347 },
348 .Float => switch (op) {
349 .Max => -std.math.inf(Child),
350 .Add => 0,
351 .Mul => 1,
352 .Min => std.math.inf(Child),
353 else => @compileError("Invalid prefixScan operation " ++ @tagName(op) ++ " for vector of floats."),
354 },
355 else => @compileError("Invalid type " ++ @typeName(VecType) ++ " for prefixScan."),
356 };
357
358 const fn_container = struct {
359 fn opFn(a: VecType, b: VecType) VecType {
360 return if (Child == bool) switch (op) {
361 .And => @select(bool, a, b, @splat(len, false)),
362 .Or => @select(bool, a, @splat(len, true), b),
363 .Xor => a != b,
364 else => unreachable,
365 } else switch (op) {
366 .And => a & b,
367 .Or => a | b,
368 .Xor => a ^ b,
369 .Add => a + b,
370 .Mul => a * b,
371 .Min => @minimum(a, b),
372 .Max => @maximum(a, b),
373 };
374 }
375 };
376
377 return prefixScanWithFunc(hop, vec, void, fn_container.opFn, identity);
378}
379
380test "vector prefix scan" {
381 if (comptime builtin.cpu.arch.isMIPS()) {
382 return error.SkipZigTest;
383 }
384
385 const int_base = Vector(4, i32){ 11, 23, 9, -21 };
386 const float_base = Vector(4, f32){ 2, 0.5, -10, 6.54321 };
387 const bool_base = Vector(4, bool){ true, false, true, false };
388
389 try std.testing.expectEqual(iota(u8, 32) + @splat(32, @as(u8, 1)), prefixScan(.Add, 1, @splat(32, @as(u8, 1))));
390 try std.testing.expectEqual(Vector(4, i32){ 11, 3, 1, 1 }, prefixScan(.And, 1, int_base));
391 try std.testing.expectEqual(Vector(4, i32){ 11, 31, 31, -1 }, prefixScan(.Or, 1, int_base));
392 try std.testing.expectEqual(Vector(4, i32){ 11, 28, 21, -2 }, prefixScan(.Xor, 1, int_base));
393 try std.testing.expectEqual(Vector(4, i32){ 11, 34, 43, 22 }, prefixScan(.Add, 1, int_base));
394 try std.testing.expectEqual(Vector(4, i32){ 11, 253, 2277, -47817 }, prefixScan(.Mul, 1, int_base));
395 try std.testing.expectEqual(Vector(4, i32){ 11, 11, 9, -21 }, prefixScan(.Min, 1, int_base));
396 try std.testing.expectEqual(Vector(4, i32){ 11, 23, 23, 23 }, prefixScan(.Max, 1, int_base));
397
398 // Trying to predict all inaccuracies when adding and multiplying floats with prefixScans would be a mess, so we don't test those.
399 try std.testing.expectEqual(Vector(4, f32){ 2, 0.5, -10, -10 }, prefixScan(.Min, 1, float_base));
400 try std.testing.expectEqual(Vector(4, f32){ 2, 2, 2, 6.54321 }, prefixScan(.Max, 1, float_base));
401
402 try std.testing.expectEqual(Vector(4, bool){ true, true, false, false }, prefixScan(.Xor, 1, bool_base));
403 try std.testing.expectEqual(Vector(4, bool){ true, true, true, true }, prefixScan(.Or, 1, bool_base));
404 try std.testing.expectEqual(Vector(4, bool){ true, false, false, false }, prefixScan(.And, 1, bool_base));
405
406 try std.testing.expectEqual(Vector(4, i32){ 11, 23, 20, 2 }, prefixScan(.Add, 2, int_base));
407 try std.testing.expectEqual(Vector(4, i32){ 22, 11, -12, -21 }, prefixScan(.Add, -1, int_base));
408 try std.testing.expectEqual(Vector(4, i32){ 11, 23, 9, -10 }, prefixScan(.Add, 3, int_base));
409}
lib/std/std.zig+1
...@@ -79,6 +79,7 @@ pub const pdb = @import("pdb.zig");...@@ -79,6 +79,7 @@ pub const pdb = @import("pdb.zig");
79pub const process = @import("process.zig");79pub const process = @import("process.zig");
80pub const rand = @import("rand.zig");80pub const rand = @import("rand.zig");
81pub const sort = @import("sort.zig");81pub const sort = @import("sort.zig");
82pub const simd = @import("simd.zig");
82pub const ascii = @import("ascii.zig");83pub const ascii = @import("ascii.zig");
83pub const testing = @import("testing.zig");84pub const testing = @import("testing.zig");
84pub const time = @import("time.zig");85pub const time = @import("time.zig");