lib/simd/src/rotate.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const builtin = @import("builtin");
2 const std = @import("std");
3
4 pub fn rotateLeft(
5 comptime D: type,
6 comptime amount: usize,
7 value: D.Vector,
8 ) D.Vector {
9 validate(D);
10 validateAmount(D, amount);
11 const amounts: UnsignedVector(D) = @splat(amount);
12 return rotateLeftBits(D, value, amounts);
13 }
14
15 pub fn rotateRight(
16 comptime D: type,
17 comptime amount: usize,
18 value: D.Vector,
19 ) D.Vector {
20 validate(D);
21 validateAmount(D, amount);
22 const amounts: UnsignedVector(D) = @splat(amount);
23 return rotateRightBits(D, value, amounts);
24 }
25
26 pub fn rotateLeftSame(comptime D: type, value: D.Vector, amount: i32) D.Vector {
27 validate(D);
28 const normalized: UnsignedLane(D) = @intCast(@mod(amount, @bitSizeOf(D.Lane)));
29 return rotateLeftBits(D, value, @splat(normalized));
30 }
31
32 pub fn rotateRightSame(comptime D: type, value: D.Vector, amount: i32) D.Vector {
33 validate(D);
34 const normalized: UnsignedLane(D) = @intCast(@mod(amount, @bitSizeOf(D.Lane)));
35 return rotateRightBits(D, value, @splat(normalized));
36 }
37
38 pub fn rol(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {
39 validate(D);
40 const raw: UnsignedVector(D) = @bitCast(amounts);
41 return rotateLeftBits(D, value, raw);
42 }
43
44 pub fn ror(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {
45 validate(D);
46 const raw: UnsignedVector(D) = @bitCast(amounts);
47 return rotateRightBits(D, value, raw);
48 }
49
50 pub fn multiRotateRight(
51 comptime D: type,
52 value: D.Vector,
53 indices: D.repartition(u8).Vector,
54 ) D.Vector {
55 validate(D);
56 if (@bitSizeOf(D.Lane) != 64) @compileError("multiRotateRight requires 64-bit lanes");
57 const bits: @Vector(D.lane_count, u64) = @bitCast(value);
58 const index_lanes: [D.lane_count * 8]u8 = indices;
59 var result: @Vector(D.lane_count, u64) = @splat(0);
60 inline for (0..D.lane_count) |lane_index| {
61 var lane: u64 = 0;
62 inline for (0..8) |byte_index| {
63 const amount = index_lanes[lane_index * 8 + byte_index] & 63;
64 const byte: u8 = @truncate(std.math.rotr(u64, bits[lane_index], amount));
65 const destination = if (builtin.cpu.arch.endian() == .little)
66 byte_index
67 else
68 byte_index ^ 7;
69 lane |= @as(u64, byte) << @intCast(destination * 8);
70 }
71 result[lane_index] = lane;
72 }
73 return @bitCast(result);
74 }
75
76 fn rotateLeftBits(
77 comptime D: type,
78 value: D.Vector,
79 raw_amounts: UnsignedVector(D),
80 ) D.Vector {
81 const bits: UnsignedVector(D) = @bitCast(value);
82 const zero: UnsignedVector(D) = @splat(0);
83 const left: AmountVector(D) = @truncate(raw_amounts);
84 const right: AmountVector(D) = @truncate(zero -% raw_amounts);
85 return @bitCast((bits << left) | (bits >> right));
86 }
87
88 fn rotateRightBits(
89 comptime D: type,
90 value: D.Vector,
91 raw_amounts: UnsignedVector(D),
92 ) D.Vector {
93 const bits: UnsignedVector(D) = @bitCast(value);
94 const zero: UnsignedVector(D) = @splat(0);
95 const right: AmountVector(D) = @truncate(raw_amounts);
96 const left: AmountVector(D) = @truncate(zero -% raw_amounts);
97 return @bitCast((bits >> right) | (bits << left));
98 }
99
100 fn UnsignedLane(comptime D: type) type {
101 return @Int(.unsigned, @bitSizeOf(D.Lane));
102 }
103
104 fn UnsignedVector(comptime D: type) type {
105 return @Vector(D.lane_count, UnsignedLane(D));
106 }
107
108 fn AmountVector(comptime D: type) type {
109 return @Vector(D.lane_count, std.math.Log2Int(UnsignedLane(D)));
110 }
111
112 fn validate(comptime D: type) void {
113 if (@typeInfo(D.Lane) != .int) @compileError("rotations require integer lanes");
114 }
115
116 fn validateAmount(comptime D: type, comptime amount: usize) void {
117 if (amount >= @bitSizeOf(D.Lane)) @compileError("rotation amount exceeds lane width");
118 }
119
120 fn verifyLaneType(comptime T: type) !void {
121 const simd = @import("root.zig");
122 const D = simd.FixedTag(T, 4);
123 const U = @Int(.unsigned, @bitSizeOf(T));
124 const UV = @Vector(4, U);
125 const raw: UV = .{ 0, 1, @as(U, 1) << (@bitSizeOf(T) - 1), std.math.maxInt(U) };
126 const value: D.Vector = @bitCast(raw);
127 try std.testing.expect(@reduce(.And, rotateRight(D, 1, rotateLeft(D, 1, value)) == value));
128 try std.testing.expect(@reduce(.And, rotateRightSame(D, rotateLeftSame(D, value, -3), -3) == value));
129 }
130
131 test "Highway rotations instantiate every integer lane type" {
132 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64 }) |T| {
133 try verifyLaneType(T);
134 }
135 }
136
137 test "Highway fixed rotations wrap every shifted bit" {
138 const simd = @import("root.zig");
139 const D = simd.FixedTag(u8, 4);
140 const value: D.Vector = .{ 0x01, 0x80, 0x55, 0x81 };
141 try std.testing.expect(@reduce(.And, rotateLeft(D, 1, value) ==
142 @as(D.Vector, .{ 0x02, 0x01, 0xaa, 0x03 })));
143 try std.testing.expect(@reduce(.And, rotateRight(D, 1, value) ==
144 @as(D.Vector, .{ 0x80, 0x40, 0xaa, 0xc0 })));
145 try std.testing.expect(@reduce(.And, rotateLeft(D, 0, value) == value));
146 }
147
148 test "Highway variable rotations mask signed amounts" {
149 const simd = @import("root.zig");
150 const D = simd.FixedTag(i16, 4);
151 const U = simd.FixedTag(u16, 4);
152 const value: D.Vector = @bitCast(@as(U.Vector, .{ 1, 0x8000, 0x1234, 0x55aa }));
153 const amounts: D.Vector = .{ 0, 1, -4, 17 };
154 const round_trip = ror(D, rol(D, value, amounts), amounts);
155 try std.testing.expect(@reduce(.And, round_trip == value));
156 try std.testing.expect(@reduce(.And, rotateLeftSame(D, value, -1) == rotateRight(D, 1, value)));
157 }
158
159 test "Highway multi rotate composes bytes within each u64 lane" {
160 const simd = @import("root.zig");
161 const D = simd.FixedTag(u64, 2);
162 const DI = D.repartition(u8);
163 const value: D.Vector = .{ 0x0102_0304_0506_0708, 0x1020_3040_5060_7080 };
164 const indices: DI.Vector = .{
165 0, 8, 16, 24, 32, 40, 48, 56,
166 56, 48, 40, 32, 24, 16, 8, 0,
167 };
168 const expected: D.Vector = if (builtin.cpu.arch.endian() == .little)
169 .{ value[0], @byteSwap(value[1]) }
170 else
171 .{ @byteSwap(value[0]), value[1] };
172 try std.testing.expect(@reduce(.And, multiRotateRight(D, value, indices) == expected));
173 }