| 1 | //! SIMD (Single Instruction; Multiple Data) convenience functions. |
| 2 | //! |
| 3 | //! May offer a potential boost in performance on some targets by performing |
| 4 | //! the same operation on multiple elements at once. |
| 5 | //! |
| 6 | //! Some functions are known to not work on MIPS. |
| 7 | |
| 8 | const std = @import("std"); |
| 9 | const builtin = @import("builtin"); |
| 10 | |
| 11 | pub fn suggestVectorLengthForCpu(comptime T: type, comptime cpu: std.Target.Cpu) ?comptime_int { |
| 12 | @setEvalBranchQuota(2_000); |
| 13 | |
| 14 | // This is guesswork, if you have better suggestions can add it or edit the current here |
| 15 | const element_bit_size = @max(8, std.math.ceilPowerOfTwo(u16, @bitSizeOf(T)) catch unreachable); |
| 16 | const vector_bit_size: u16 = blk: { |
| 17 | if (cpu.arch.isX86()) { |
| 18 | if (T == bool and cpu.has(.x86, .prefer_mask_registers)) return 64; |
| 19 | if (builtin.zig_backend != .stage2_x86_64 and cpu.has(.x86, .avx512f) and !cpu.hasAny(.x86, &.{ .prefer_256_bit, .prefer_128_bit })) break :blk 512; |
| 20 | if (cpu.hasAny(.x86, &.{ .prefer_256_bit, .avx2 }) and !cpu.has(.x86, .prefer_128_bit)) break :blk 256; |
| 21 | if (cpu.has(.x86, .sse)) break :blk 128; |
| 22 | if (cpu.hasAny(.x86, &.{ .mmx, .@"3dnow" })) break :blk 64; |
| 23 | } else if (cpu.arch.isArm()) { |
| 24 | if (cpu.has(.arm, .neon)) break :blk 128; |
| 25 | } else if (cpu.arch.isAARCH64()) { |
| 26 | // NVIDIA Grace supports 128-bit SVE |
| 27 | // AWS Graviton3 supports 256-bit SVE |
| 28 | // Fujitsu A64FX supports 512-bit SVE |
| 29 | // -> 256-bit seems like a good default for now. |
| 30 | if (cpu.has(.aarch64, .sve)) break :blk 256; |
| 31 | if (cpu.has(.aarch64, .neon)) break :blk 128; |
| 32 | } else if (cpu.arch == .hexagon) { |
| 33 | if (cpu.has(.hexagon, .hvx_length64b)) break :blk 512; |
| 34 | if (cpu.has(.hexagon, .hvx)) break :blk 1024; |
| 35 | } else if (cpu.arch.isLoongArch()) { |
| 36 | if (cpu.has(.loongarch, .lasx)) break :blk 256; |
| 37 | if (cpu.has(.loongarch, .lsx)) break :blk 128; |
| 38 | } else if (cpu.arch.isMIPS()) { |
| 39 | if (cpu.has(.mips, .msa)) break :blk 128; |
| 40 | if (cpu.has(.mips, .mips3d)) break :blk 64; |
| 41 | } else if (cpu.arch.isPowerPC()) { |
| 42 | if (cpu.has(.powerpc, .vsx)) break :blk 128; |
| 43 | if (cpu.has(.powerpc, .altivec)) break :blk 128; |
| 44 | } else if (cpu.arch.isRISCV()) { |
| 45 | // In RISC-V Vector Registers are length agnostic so there's no good way to determine the best size. |
| 46 | // The usual vector length in most RISC-V cpus is 256 bits, however it can get to multiple kB. |
| 47 | if (cpu.has(.riscv, .v)) { |
| 48 | inline for (.{ |
| 49 | .{ .zvl65536b, 65536 }, |
| 50 | .{ .zvl32768b, 32768 }, |
| 51 | .{ .zvl16384b, 16384 }, |
| 52 | .{ .zvl8192b, 8192 }, |
| 53 | .{ .zvl4096b, 4096 }, |
| 54 | .{ .zvl2048b, 2048 }, |
| 55 | .{ .zvl1024b, 1024 }, |
| 56 | .{ .zvl512b, 512 }, |
| 57 | .{ .zvl256b, 256 }, |
| 58 | .{ .zvl128b, 128 }, |
| 59 | .{ .zvl64b, 64 }, |
| 60 | .{ .zvl32b, 32 }, |
| 61 | }) |mapping| { |
| 62 | if (cpu.has(.riscv, mapping[0])) break :blk mapping[1]; |
| 63 | } |
| 64 | |
| 65 | break :blk 256; |
| 66 | } |
| 67 | } else if (cpu.arch == .s390x) { |
| 68 | if (cpu.has(.s390x, .vector)) break :blk 128; |
| 69 | } else if (cpu.arch.isSPARC()) { |
| 70 | if (cpu.hasAny(.sparc, &.{ .vis, .vis2, .vis3 })) break :blk 64; |
| 71 | } else if (cpu.arch == .kvx) { |
| 72 | break :blk 1024; |
| 73 | } else if (cpu.arch == .ve) { |
| 74 | if (cpu.has(.ve, .vpu)) break :blk 2048; |
| 75 | } else if (cpu.arch.isWasm()) { |
| 76 | if (cpu.has(.wasm, .simd128)) break :blk 128; |
| 77 | } |
| 78 | return null; |
| 79 | }; |
| 80 | if (vector_bit_size <= element_bit_size) return null; |
| 81 | |
| 82 | return @divExact(vector_bit_size, element_bit_size); |
| 83 | } |
| 84 | |
| 85 | /// Suggests a target-dependant vector length for a given type, or null if scalars are recommended. |
| 86 | /// Not yet implemented for every CPU architecture. |
| 87 | pub fn suggestVectorLength(comptime T: type) ?comptime_int { |
| 88 | return suggestVectorLengthForCpu(T, builtin.cpu); |
| 89 | } |
| 90 | |
| 91 | test "suggestVectorLengthForCpu works with signed and unsigned values" { |
| 92 | comptime var cpu = std.Target.Cpu.baseline(std.Target.Cpu.Arch.x86_64, builtin.os); |
| 93 | comptime cpu.features.addFeature(@backingInt(std.Target.x86.Feature.avx512f)); |
| 94 | comptime cpu.features.populateDependencies(&std.Target.x86.all_features); |
| 95 | const expected_len: usize = switch (builtin.zig_backend) { |
| 96 | .stage2_x86_64 => 8, |
| 97 | else => 16, |
| 98 | }; |
| 99 | const signed_integer_len = suggestVectorLengthForCpu(i32, cpu).?; |
| 100 | const unsigned_integer_len = suggestVectorLengthForCpu(u32, cpu).?; |
| 101 | try std.testing.expectEqual(expected_len, unsigned_integer_len); |
| 102 | try std.testing.expectEqual(expected_len, signed_integer_len); |
| 103 | } |
| 104 | |
| 105 | fn vectorLength(comptime VectorType: type) comptime_int { |
| 106 | return switch (@typeInfo(VectorType)) { |
| 107 | .vector => |info| info.len, |
| 108 | .array => |info| info.len, |
| 109 | else => @compileError("Invalid type " ++ @typeName(VectorType)), |
| 110 | }; |
| 111 | } |
| 112 | |
| 113 | /// Returns the smallest type of unsigned ints capable of indexing any element within the given vector type. |
| 114 | pub fn VectorIndex(comptime VectorType: type) type { |
| 115 | return std.math.IntFittingRange(0, vectorLength(VectorType) - 1); |
| 116 | } |
| 117 | |
| 118 | /// Returns the smallest type of unsigned ints capable of holding the length of the given vector type. |
| 119 | pub fn VectorCount(comptime VectorType: type) type { |
| 120 | return std.math.IntFittingRange(0, vectorLength(VectorType)); |
| 121 | } |
| 122 | |
| 123 | /// Returns a vector containing the first `len` integers in order from 0 to `len`-1. |
| 124 | /// For example, `iota(i32, 8)` will return a vector containing `.{0, 1, 2, 3, 4, 5, 6, 7}`. |
| 125 | pub inline fn iota(comptime T: type, comptime len: usize) @Vector(len, T) { |
| 126 | comptime { |
| 127 | var out: [len]T = undefined; |
| 128 | for (&out, 0..) |*element, i| { |
| 129 | element.* = switch (@typeInfo(T)) { |
| 130 | .int => @as(T, @intCast(i)), |
| 131 | .float => @as(T, @floatFromInt(i)), |
| 132 | else => @compileError("Can't use type " ++ @typeName(T) ++ " in iota."), |
| 133 | }; |
| 134 | } |
| 135 | return @as(@Vector(len, T), out); |
| 136 | } |
| 137 | } |
| 138 | |
| 139 | /// Returns a vector containing the same elements as the input, but repeated until the desired length is reached. |
| 140 | /// For example, `repeat(8, [_]u32{1, 2, 3})` will return a vector containing `.{1, 2, 3, 1, 2, 3, 1, 2}`. |
| 141 | pub fn repeat(comptime len: usize, vec: anytype) @Vector(len, std.meta.Child(@TypeOf(vec))) { |
| 142 | const Child = std.meta.Child(@TypeOf(vec)); |
| 143 | |
| 144 | return @shuffle(Child, vec, undefined, iota(i32, len) % @as(@Vector(len, i32), @splat(@intCast(vectorLength(@TypeOf(vec)))))); |
| 145 | } |
| 146 | |
| 147 | /// Returns a vector containing all elements of the first vector at the lower indices followed by all elements of the second vector |
| 148 | /// at the higher indices. |
| 149 | pub fn join(a: anytype, b: anytype) @Vector(vectorLength(@TypeOf(a)) + vectorLength(@TypeOf(b)), std.meta.Child(@TypeOf(a))) { |
| 150 | const Child = std.meta.Child(@TypeOf(a)); |
| 151 | const a_len = vectorLength(@TypeOf(a)); |
| 152 | const b_len = vectorLength(@TypeOf(b)); |
| 153 | |
| 154 | return @shuffle(Child, a, b, @as([a_len]i32, iota(i32, a_len)) ++ @as([b_len]i32, ~iota(i32, b_len))); |
| 155 | } |
| 156 | |
| 157 | /// Returns a vector whose elements alternates between those of each input vector. |
| 158 | /// 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}`. |
| 159 | pub fn interlace(vecs: anytype) @Vector(vectorLength(@TypeOf(vecs[0])) * vecs.len, std.meta.Child(@TypeOf(vecs[0]))) { |
| 160 | const VecType = @TypeOf(vecs[0]); |
| 161 | const vecs_arr = @as([vecs.len]VecType, vecs); |
| 162 | const Child = std.meta.Child(@TypeOf(vecs_arr[0])); |
| 163 | |
| 164 | if (vecs_arr.len == 1) return vecs_arr[0]; |
| 165 | |
| 166 | const a_vec_count = (1 + vecs_arr.len) >> 1; |
| 167 | const b_vec_count = vecs_arr.len >> 1; |
| 168 | |
| 169 | const a = interlace(@as(*const [a_vec_count]VecType, @ptrCast(vecs_arr[0..a_vec_count])).*); |
| 170 | const b = interlace(@as(*const [b_vec_count]VecType, @ptrCast(vecs_arr[a_vec_count..])).*); |
| 171 | |
| 172 | const a_len = vectorLength(@TypeOf(a)); |
| 173 | const b_len = vectorLength(@TypeOf(b)); |
| 174 | const len = a_len + b_len; |
| 175 | |
| 176 | const indices = comptime blk: { |
| 177 | const Vi32 = @Vector(len, i32); |
| 178 | const count_up = iota(i32, len); |
| 179 | const cycle = @divFloor(count_up, @as(Vi32, @splat(@intCast(vecs_arr.len)))); |
| 180 | const select_mask = repeat(len, join(@as(@Vector(a_vec_count, bool), @splat(true)), @as(@Vector(b_vec_count, bool), @splat(false)))); |
| 181 | const a_indices = count_up - cycle * @as(Vi32, @splat(@intCast(b_vec_count))); |
| 182 | const b_indices = shiftElementsRight(count_up - cycle * @as(Vi32, @splat(@intCast(a_vec_count))), a_vec_count, 0); |
| 183 | break :blk @select(i32, select_mask, a_indices, ~b_indices); |
| 184 | }; |
| 185 | |
| 186 | return @shuffle(Child, a, b, indices); |
| 187 | } |
| 188 | |
| 189 | /// The contents of `interlaced` is evenly split between vec_count vectors that are returned as an array. They "take turns", |
| 190 | /// receiving one element from `interlaced` at a time. |
| 191 | pub fn deinterlace( |
| 192 | comptime vec_count: usize, |
| 193 | interlaced: anytype, |
| 194 | ) [vec_count]@Vector( |
| 195 | vectorLength(@TypeOf(interlaced)) / vec_count, |
| 196 | std.meta.Child(@TypeOf(interlaced)), |
| 197 | ) { |
| 198 | const vec_len = vectorLength(@TypeOf(interlaced)) / vec_count; |
| 199 | const Child = std.meta.Child(@TypeOf(interlaced)); |
| 200 | |
| 201 | var out: [vec_count]@Vector(vec_len, Child) = undefined; |
| 202 | |
| 203 | comptime var i: usize = 0; // for-loops don't work for this, apparently. |
| 204 | inline while (i < out.len) : (i += 1) { |
| 205 | const indices = comptime iota(i32, vec_len) * @as(@Vector(vec_len, i32), @splat(@intCast(vec_count))) + @as(@Vector(vec_len, i32), @splat(@intCast(i))); |
| 206 | out[i] = @shuffle(Child, interlaced, undefined, indices); |
| 207 | } |
| 208 | |
| 209 | return out; |
| 210 | } |
| 211 | |
| 212 | pub fn extract( |
| 213 | vec: anytype, |
| 214 | comptime first: VectorIndex(@TypeOf(vec)), |
| 215 | comptime count: VectorCount(@TypeOf(vec)), |
| 216 | ) @Vector(count, std.meta.Child(@TypeOf(vec))) { |
| 217 | const Child = std.meta.Child(@TypeOf(vec)); |
| 218 | const len = vectorLength(@TypeOf(vec)); |
| 219 | |
| 220 | std.debug.assert(@as(comptime_int, @intCast(first)) + @as(comptime_int, @intCast(count)) <= len); |
| 221 | |
| 222 | return @shuffle(Child, vec, undefined, iota(i32, count) + @as(@Vector(count, i32), @splat(@intCast(first)))); |
| 223 | } |
| 224 | |
| 225 | test "vector patterns" { |
| 226 | if (builtin.cpu.arch == .hexagon) return error.SkipZigTest; |
| 227 | |
| 228 | const base = @Vector(4, u32){ 10, 20, 30, 40 }; |
| 229 | const other_base = @Vector(4, u32){ 55, 66, 77, 88 }; |
| 230 | |
| 231 | const small_bases = [5]@Vector(2, u8){ |
| 232 | @Vector(2, u8){ 0, 1 }, |
| 233 | @Vector(2, u8){ 2, 3 }, |
| 234 | @Vector(2, u8){ 4, 5 }, |
| 235 | @Vector(2, u8){ 6, 7 }, |
| 236 | @Vector(2, u8){ 8, 9 }, |
| 237 | }; |
| 238 | |
| 239 | try std.testing.expectEqual([6]u32{ 10, 20, 30, 40, 10, 20 }, repeat(6, base)); |
| 240 | try std.testing.expectEqual([8]u32{ 10, 20, 30, 40, 55, 66, 77, 88 }, join(base, other_base)); |
| 241 | try std.testing.expectEqual([2]u32{ 20, 30 }, extract(base, 1, 2)); |
| 242 | |
| 243 | try std.testing.expectEqual([8]u32{ 10, 55, 20, 66, 30, 77, 40, 88 }, interlace(.{ base, other_base })); |
| 244 | |
| 245 | const small_braid = interlace(small_bases); |
| 246 | try std.testing.expectEqual([10]u8{ 0, 2, 4, 6, 8, 1, 3, 5, 7, 9 }, small_braid); |
| 247 | try std.testing.expectEqual(small_bases, deinterlace(small_bases.len, small_braid)); |
| 248 | } |
| 249 | |
| 250 | /// Joins two vectors, shifts them leftwards (towards lower indices) and extracts the leftmost elements into a vector the length of a and b. |
| 251 | pub fn mergeShift(a: anytype, b: anytype, comptime shift: VectorCount(@TypeOf(a, b))) @TypeOf(a, b) { |
| 252 | const len = vectorLength(@TypeOf(a, b)); |
| 253 | |
| 254 | return extract(join(a, b), shift, len); |
| 255 | } |
| 256 | |
| 257 | /// Elements are shifted rightwards (towards higher indices). New elements are added to the left, and the rightmost elements are cut off |
| 258 | /// so that the length of the vector stays the same. |
| 259 | pub fn shiftElementsRight(vec: anytype, comptime amount: VectorCount(@TypeOf(vec)), shift_in: std.meta.Child(@TypeOf(vec))) @TypeOf(vec) { |
| 260 | // It may be possible to implement shifts and rotates with a runtime-friendly slice of two joined vectors, as the length of the |
| 261 | // slice would be comptime-known. This would permit vector shifts and rotates by a non-comptime-known amount. |
| 262 | // However, I am unsure whether compiler optimizations would handle that well enough on all platforms. |
| 263 | const V = @TypeOf(vec); |
| 264 | const len = vectorLength(V); |
| 265 | |
| 266 | return mergeShift(@as(V, @splat(shift_in)), vec, len - amount); |
| 267 | } |
| 268 | |
| 269 | /// Elements are shifted leftwards (towards lower indices). New elements are added to the right, and the leftmost elements are cut off |
| 270 | /// so that no elements with indices below 0 remain. |
| 271 | pub fn shiftElementsLeft(vec: anytype, comptime amount: VectorCount(@TypeOf(vec)), shift_in: std.meta.Child(@TypeOf(vec))) @TypeOf(vec) { |
| 272 | const V = @TypeOf(vec); |
| 273 | |
| 274 | return mergeShift(vec, @as(V, @splat(shift_in)), amount); |
| 275 | } |
| 276 | |
| 277 | /// Elements are shifted leftwards (towards lower indices). Elements that leave to the left will reappear to the right in the same order. |
| 278 | pub fn rotateElementsLeft(vec: anytype, comptime amount: VectorCount(@TypeOf(vec))) @TypeOf(vec) { |
| 279 | return mergeShift(vec, vec, amount); |
| 280 | } |
| 281 | |
| 282 | /// Elements are shifted rightwards (towards higher indices). Elements that leave to the right will reappear to the left in the same order. |
| 283 | pub fn rotateElementsRight(vec: anytype, comptime amount: VectorCount(@TypeOf(vec))) @TypeOf(vec) { |
| 284 | return rotateElementsLeft(vec, vectorLength(@TypeOf(vec)) - amount); |
| 285 | } |
| 286 | |
| 287 | pub fn reverseOrder(vec: anytype) @TypeOf(vec) { |
| 288 | const Child = std.meta.Child(@TypeOf(vec)); |
| 289 | const len = vectorLength(@TypeOf(vec)); |
| 290 | |
| 291 | return @shuffle(Child, vec, undefined, @as(@Vector(len, i32), @splat(@as(i32, @intCast(len)) - 1)) - iota(i32, len)); |
| 292 | } |
| 293 | |
| 294 | test "vector shifting" { |
| 295 | const base = @Vector(4, u32){ 10, 20, 30, 40 }; |
| 296 | |
| 297 | try std.testing.expectEqual([4]u32{ 30, 40, 999, 999 }, shiftElementsLeft(base, 2, 999)); |
| 298 | try std.testing.expectEqual([4]u32{ 999, 999, 10, 20 }, shiftElementsRight(base, 2, 999)); |
| 299 | try std.testing.expectEqual([4]u32{ 20, 30, 40, 10 }, rotateElementsLeft(base, 1)); |
| 300 | try std.testing.expectEqual([4]u32{ 40, 10, 20, 30 }, rotateElementsRight(base, 1)); |
| 301 | try std.testing.expectEqual([4]u32{ 40, 30, 20, 10 }, reverseOrder(base)); |
| 302 | } |
| 303 | |
| 304 | pub fn firstTrue(vec: anytype) ?VectorIndex(@TypeOf(vec)) { |
| 305 | const len = vectorLength(@TypeOf(vec)); |
| 306 | const IndexInt = VectorIndex(@TypeOf(vec)); |
| 307 | |
| 308 | if (!@reduce(.Or, vec)) { |
| 309 | return null; |
| 310 | } |
| 311 | const all_max: @Vector(len, IndexInt) = @splat(~@as(IndexInt, 0)); |
| 312 | const indices = @select(IndexInt, vec, iota(IndexInt, len), all_max); |
| 313 | return @reduce(.Min, indices); |
| 314 | } |
| 315 | |
| 316 | pub fn lastTrue(vec: anytype) ?VectorIndex(@TypeOf(vec)) { |
| 317 | const len = vectorLength(@TypeOf(vec)); |
| 318 | const IndexInt = VectorIndex(@TypeOf(vec)); |
| 319 | |
| 320 | if (!@reduce(.Or, vec)) { |
| 321 | return null; |
| 322 | } |
| 323 | |
| 324 | const all_zeroes: @Vector(len, IndexInt) = @splat(0); |
| 325 | const indices = @select(IndexInt, vec, iota(IndexInt, len), all_zeroes); |
| 326 | return @reduce(.Max, indices); |
| 327 | } |
| 328 | |
| 329 | pub fn countTrues(vec: anytype) VectorCount(@TypeOf(vec)) { |
| 330 | const len = vectorLength(@TypeOf(vec)); |
| 331 | const CountIntType = VectorCount(@TypeOf(vec)); |
| 332 | |
| 333 | const all_ones: @Vector(len, CountIntType) = @splat(1); |
| 334 | const all_zeroes: @Vector(len, CountIntType) = @splat(0); |
| 335 | |
| 336 | const one_if_true = @select(CountIntType, vec, all_ones, all_zeroes); |
| 337 | return @reduce(.Add, one_if_true); |
| 338 | } |
| 339 | |
| 340 | pub fn firstIndexOfValue(vec: anytype, value: std.meta.Child(@TypeOf(vec))) ?VectorIndex(@TypeOf(vec)) { |
| 341 | const V = @TypeOf(vec); |
| 342 | |
| 343 | return firstTrue(vec == @as(V, @splat(value))); |
| 344 | } |
| 345 | |
| 346 | pub fn lastIndexOfValue(vec: anytype, value: std.meta.Child(@TypeOf(vec))) ?VectorIndex(@TypeOf(vec)) { |
| 347 | const V = @TypeOf(vec); |
| 348 | |
| 349 | return lastTrue(vec == @as(V, @splat(value))); |
| 350 | } |
| 351 | |
| 352 | pub fn countElementsWithValue(vec: anytype, value: std.meta.Child(@TypeOf(vec))) VectorCount(@TypeOf(vec)) { |
| 353 | const V = @TypeOf(vec); |
| 354 | |
| 355 | return countTrues(vec == @as(V, @splat(value))); |
| 356 | } |
| 357 | |
| 358 | test "vector searching" { |
| 359 | const base = @Vector(8, u32){ 6, 4, 7, 4, 4, 2, 3, 7 }; |
| 360 | |
| 361 | try std.testing.expectEqual(@as(?u3, 1), firstIndexOfValue(base, 4)); |
| 362 | try std.testing.expectEqual(@as(?u3, 4), lastIndexOfValue(base, 4)); |
| 363 | try std.testing.expectEqual(@as(?u3, null), lastIndexOfValue(base, 99)); |
| 364 | try std.testing.expectEqual(@as(u4, 3), countElementsWithValue(base, 4)); |
| 365 | } |
| 366 | |
| 367 | /// Same as prefixScan, but with a user-provided, mathematically associative function. |
| 368 | pub fn prefixScanWithFunc( |
| 369 | comptime hop: isize, |
| 370 | vec: anytype, |
| 371 | /// The error type that `func` might return. Set this to `void` if `func` doesn't return an error union. |
| 372 | comptime ErrorType: type, |
| 373 | comptime func: fn (@TypeOf(vec), @TypeOf(vec)) if (ErrorType == void) @TypeOf(vec) else ErrorType!@TypeOf(vec), |
| 374 | /// When one operand of the operation performed by `func` is this value, the result must equal the other operand. |
| 375 | /// For example, this should be 0 for addition or 1 for multiplication. |
| 376 | comptime identity: std.meta.Child(@TypeOf(vec)), |
| 377 | ) if (ErrorType == void) @TypeOf(vec) else ErrorType!@TypeOf(vec) { |
| 378 | const len = vectorLength(@TypeOf(vec)); |
| 379 | |
| 380 | if (hop == 0) @compileError("hop can not be 0; you'd be going nowhere forever!"); |
| 381 | const abs_hop = if (hop < 0) -hop else hop; |
| 382 | |
| 383 | var acc = vec; |
| 384 | comptime var i = 0; |
| 385 | inline while ((abs_hop << i) < len) : (i += 1) { |
| 386 | const shifted = if (hop < 0) shiftElementsLeft(acc, abs_hop << i, identity) else shiftElementsRight(acc, abs_hop << i, identity); |
| 387 | |
| 388 | acc = if (ErrorType == void) func(acc, shifted) else try func(acc, shifted); |
| 389 | } |
| 390 | return acc; |
| 391 | } |
| 392 | |
| 393 | /// Returns a vector whose elements are the result of performing the specified operation on the corresponding |
| 394 | /// element of the input vector and every hop'th element that came before it (or after, if hop is negative). |
| 395 | /// Supports the same operations as the @reduce() builtin. Takes O(logN) to compute. |
| 396 | /// The scan is not linear, which may affect floating point errors. This may affect the determinism of |
| 397 | /// algorithms that use this function. |
| 398 | pub fn prefixScan(comptime op: std.builtin.ReduceOp, comptime hop: isize, vec: anytype) @TypeOf(vec) { |
| 399 | const VecType = @TypeOf(vec); |
| 400 | const Child = std.meta.Child(VecType); |
| 401 | |
| 402 | const identity = comptime switch (@typeInfo(Child)) { |
| 403 | .bool => switch (op) { |
| 404 | .Or, .Xor => false, |
| 405 | .And => true, |
| 406 | else => @compileError("Invalid prefixScan operation " ++ @tagName(op) ++ " for vector of booleans."), |
| 407 | }, |
| 408 | .int => switch (op) { |
| 409 | .Max => std.math.minInt(Child), |
| 410 | .Add, .Or, .Xor => 0, |
| 411 | .Mul => 1, |
| 412 | .And, .Min => std.math.maxInt(Child), |
| 413 | }, |
| 414 | .float => switch (op) { |
| 415 | .Max => -std.math.inf(Child), |
| 416 | .Add => 0, |
| 417 | .Mul => 1, |
| 418 | .Min => std.math.inf(Child), |
| 419 | else => @compileError("Invalid prefixScan operation " ++ @tagName(op) ++ " for vector of floats."), |
| 420 | }, |
| 421 | else => @compileError("Invalid type " ++ @typeName(VecType) ++ " for prefixScan."), |
| 422 | }; |
| 423 | |
| 424 | const fn_container = struct { |
| 425 | fn opFn(a: VecType, b: VecType) VecType { |
| 426 | return if (Child == bool) switch (op) { |
| 427 | .And => @select(bool, a, b, @as(VecType, @splat(false))), |
| 428 | .Or => @select(bool, a, @as(VecType, @splat(true)), b), |
| 429 | .Xor => a != b, |
| 430 | else => unreachable, |
| 431 | } else switch (op) { |
| 432 | .And => a & b, |
| 433 | .Or => a | b, |
| 434 | .Xor => a ^ b, |
| 435 | .Add => a + b, |
| 436 | .Mul => a * b, |
| 437 | .Min => @min(a, b), |
| 438 | .Max => @max(a, b), |
| 439 | }; |
| 440 | } |
| 441 | }; |
| 442 | |
| 443 | return prefixScanWithFunc(hop, vec, void, fn_container.opFn, identity); |
| 444 | } |
| 445 | |
| 446 | test "vector prefix scan" { |
| 447 | const int_base = @Vector(4, i32){ 11, 23, 9, -21 }; |
| 448 | const float_base = @Vector(4, f32){ 2, 0.5, -10, 6.54321 }; |
| 449 | const bool_base = @Vector(4, bool){ true, false, true, false }; |
| 450 | |
| 451 | const ones: @Vector(32, u8) = @splat(1); |
| 452 | |
| 453 | try std.testing.expectEqual(iota(u8, 32) + ones, prefixScan(.Add, 1, ones)); |
| 454 | try std.testing.expectEqual(@Vector(4, i32){ 11, 3, 1, 1 }, prefixScan(.And, 1, int_base)); |
| 455 | try std.testing.expectEqual(@Vector(4, i32){ 11, 31, 31, -1 }, prefixScan(.Or, 1, int_base)); |
| 456 | try std.testing.expectEqual(@Vector(4, i32){ 11, 28, 21, -2 }, prefixScan(.Xor, 1, int_base)); |
| 457 | try std.testing.expectEqual(@Vector(4, i32){ 11, 34, 43, 22 }, prefixScan(.Add, 1, int_base)); |
| 458 | try std.testing.expectEqual(@Vector(4, i32){ 11, 253, 2277, -47817 }, prefixScan(.Mul, 1, int_base)); |
| 459 | try std.testing.expectEqual(@Vector(4, i32){ 11, 11, 9, -21 }, prefixScan(.Min, 1, int_base)); |
| 460 | try std.testing.expectEqual(@Vector(4, i32){ 11, 23, 23, 23 }, prefixScan(.Max, 1, int_base)); |
| 461 | |
| 462 | // Trying to predict all inaccuracies when adding and multiplying floats with prefixScans would be a mess, so we don't test those. |
| 463 | try std.testing.expectEqual(@Vector(4, f32){ 2, 0.5, -10, -10 }, prefixScan(.Min, 1, float_base)); |
| 464 | try std.testing.expectEqual(@Vector(4, f32){ 2, 2, 2, 6.54321 }, prefixScan(.Max, 1, float_base)); |
| 465 | |
| 466 | try std.testing.expectEqual(@Vector(4, bool){ true, true, false, false }, prefixScan(.Xor, 1, bool_base)); |
| 467 | try std.testing.expectEqual(@Vector(4, bool){ true, true, true, true }, prefixScan(.Or, 1, bool_base)); |
| 468 | try std.testing.expectEqual(@Vector(4, bool){ true, false, false, false }, prefixScan(.And, 1, bool_base)); |
| 469 | |
| 470 | try std.testing.expectEqual(@Vector(4, i32){ 11, 23, 20, 2 }, prefixScan(.Add, 2, int_base)); |
| 471 | try std.testing.expectEqual(@Vector(4, i32){ 22, 11, -12, -21 }, prefixScan(.Add, -1, int_base)); |
| 472 | try std.testing.expectEqual(@Vector(4, i32){ 11, 23, 9, -10 }, prefixScan(.Add, 3, int_base)); |
| 473 | } |