lib/simd/src/swizzle.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn extractLane(comptime D: type, value: D.Vector, index: usize) D.Lane {
4 std.debug.assert(index < D.lane_count);
5 const lanes: [D.lane_count]D.Lane = @bitCast(value);
6 return lanes[index];
7 }
8
9 pub fn broadcastLane(comptime D: type, comptime index: usize, value: D.Vector) D.Vector {
10 if (index >= D.lane_count) @compileError("broadcast lane is outside the vector");
11 return @splat(value[index]);
12 }
13
14 pub fn dupEven(comptime D: type, value: D.Vector) D.Vector {
15 var result: D.Vector = undefined;
16 inline for (0..D.lane_count) |index| result[index] = value[index & ~@as(usize, 1)];
17 return result;
18 }
19
20 pub fn dupOdd(comptime D: type, value: D.Vector) D.Vector {
21 if (D.lane_count < 2) @compileError("dupOdd requires at least two lanes");
22 var result: D.Vector = undefined;
23 inline for (0..D.lane_count) |index| result[index] = value[(index & ~@as(usize, 1)) + 1];
24 return result;
25 }
26
27 pub fn oddEven(comptime D: type, odd: D.Vector, even: D.Vector) D.Vector {
28 var result: D.Vector = undefined;
29 inline for (0..D.lane_count) |index| {
30 result[index] = if (index & 1 == 0) even[index] else odd[index];
31 }
32 return result;
33 }
34
35 pub fn reverse(comptime D: type, value: D.Vector) D.Vector {
36 var result: D.Vector = undefined;
37 inline for (0..D.lane_count) |index| result[index] = value[D.lane_count - 1 - index];
38 return result;
39 }
40
41 pub fn reverse2(comptime D: type, value: D.Vector) D.Vector {
42 if (D.lane_count == 1) return value;
43 return reverseGroups(D, 2, value);
44 }
45
46 pub fn reverse4(comptime D: type, value: D.Vector) D.Vector {
47 return reverseGroups(D, 4, value);
48 }
49
50 pub fn reverse8(comptime D: type, value: D.Vector) D.Vector {
51 return reverseGroups(D, 8, value);
52 }
53
54 pub fn reverseLaneBytes(comptime D: type, value: D.Vector) D.Vector {
55 if (@typeInfo(D.Lane) != .int) @compileError("reverseLaneBytes requires integer lanes");
56 var result: D.Vector = undefined;
57 inline for (0..D.lane_count) |index| result[index] = @byteSwap(value[index]);
58 return result;
59 }
60
61 pub fn reverseBits(comptime D: type, value: D.Vector) D.Vector {
62 if (@typeInfo(D.Lane) != .int) @compileError("reverseBits requires integer lanes");
63 var result: D.Vector = undefined;
64 inline for (0..D.lane_count) |index| result[index] = @bitReverse(value[index]);
65 return result;
66 }
67
68 pub fn reverseBlocks(comptime D: type, value: D.Vector) D.Vector {
69 if (D.byte_count < 16 or D.byte_count % 16 != 0) {
70 @compileError("reverseBlocks requires complete 128-bit blocks");
71 }
72 const lanes_per_block = 16 / @sizeOf(D.Lane);
73 const block_count = D.lane_count / lanes_per_block;
74 var result: D.Vector = undefined;
75 inline for (0..D.lane_count) |index| {
76 const block = index / lanes_per_block;
77 const lane = index % lanes_per_block;
78 result[index] = value[(block_count - 1 - block) * lanes_per_block + lane];
79 }
80 return result;
81 }
82
83 fn reverseGroups(comptime D: type, comptime group: usize, value: D.Vector) D.Vector {
84 if (D.lane_count < group or D.lane_count % group != 0) {
85 @compileError("lane count must contain complete reversal groups");
86 }
87 var result: D.Vector = undefined;
88 inline for (0..D.lane_count) |index| result[index] = value[index ^ (group - 1)];
89 return result;
90 }
91
92 fn verifyLaneType(comptime T: type) !void {
93 const simd = @import("root.zig");
94 const D = simd.FixedTag(T, 8);
95 const value: D.Vector = @splat(0);
96 try std.testing.expect(@reduce(.And, reverse(D, reverse(D, value)) == value));
97 try std.testing.expect(@reduce(.And, dupEven(D, value) == value));
98 try std.testing.expect(@reduce(.And, dupOdd(D, value) == value));
99 try std.testing.expect(@reduce(.And, oddEven(D, value, value) == value));
100 }
101
102 test "Highway whole-vector swizzles instantiate every lane type" {
103 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {
104 try verifyLaneType(T);
105 }
106 }
107
108 test "Highway lane duplication selection and broadcast retain lane order" {
109 const simd = @import("root.zig");
110 const D = simd.FixedTag(i32, 8);
111 const value: D.Vector = .{ 1, 2, 3, 4, 5, 6, 7, 8 };
112 try std.testing.expectEqual(@as(i32, 5), extractLane(D, value, 4));
113 try std.testing.expect(@reduce(.And, broadcastLane(D, 4, value) == @as(D.Vector, @splat(5))));
114 try std.testing.expect(@reduce(.And, dupEven(D, value) == @as(D.Vector, .{ 1, 1, 3, 3, 5, 5, 7, 7 })));
115 try std.testing.expect(@reduce(.And, dupOdd(D, value) == @as(D.Vector, .{ 2, 2, 4, 4, 6, 6, 8, 8 })));
116 try std.testing.expect(@reduce(.And, oddEven(D, value + @as(D.Vector, @splat(10)), value) ==
117 @as(D.Vector, .{ 1, 12, 3, 14, 5, 16, 7, 18 })));
118 }
119
120 test "Highway reversals distinguish whole vectors groups lanes bits and blocks" {
121 const simd = @import("root.zig");
122 const D = simd.FixedTag(u16, 8);
123 const value: D.Vector = .{ 0x0001, 0x0002, 0x0004, 0x0008, 0x0010, 0x0020, 0x0040, 0x0080 };
124 try std.testing.expect(@reduce(.And, reverse(D, value) == @as(D.Vector, .{
125 0x0080, 0x0040, 0x0020, 0x0010, 0x0008, 0x0004, 0x0002, 0x0001,
126 })));
127 try std.testing.expect(@reduce(.And, reverse2(D, value) == @as(D.Vector, .{
128 0x0002, 0x0001, 0x0008, 0x0004, 0x0020, 0x0010, 0x0080, 0x0040,
129 })));
130 try std.testing.expect(@reduce(.And, reverse4(D, value) == @as(D.Vector, .{
131 0x0008, 0x0004, 0x0002, 0x0001, 0x0080, 0x0040, 0x0020, 0x0010,
132 })));
133 try std.testing.expect(@reduce(.And, reverse8(D, value) == reverse(D, value)));
134 try std.testing.expect(@reduce(.And, reverseLaneBytes(D, value) == @as(D.Vector, .{
135 0x0100, 0x0200, 0x0400, 0x0800, 0x1000, 0x2000, 0x4000, 0x8000,
136 })));
137 try std.testing.expect(@reduce(.And, reverseBits(D, value) == @as(D.Vector, .{
138 0x8000, 0x4000, 0x2000, 0x1000, 0x0800, 0x0400, 0x0200, 0x0100,
139 })));
140
141 const B = simd.FixedTag(u32, 8);
142 try std.testing.expect(@reduce(.And, reverseBlocks(B, @as(B.Vector, .{ 0, 1, 2, 3, 4, 5, 6, 7 })) ==
143 @as(B.Vector, .{ 4, 5, 6, 7, 0, 1, 2, 3 })));
144 }