lib/simd/src/memory.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn load(comptime D: type, input: []const D.Lane) D.Vector {
4 std.debug.assert(input.len >= D.lane_count);
5 return @as(D.Vector, input[0..D.lane_count].*);
6 }
7
8 pub fn store(comptime D: type, value: D.Vector, output: []D.Lane) void {
9 std.debug.assert(output.len >= D.lane_count);
10 output[0..D.lane_count].* = value;
11 }
12
13 pub fn loadN(comptime D: type, input: []const D.Lane, count: usize) D.Vector {
14 const lanes = @min(count, D.lane_count);
15 std.debug.assert(input.len >= lanes);
16 var result: D.Vector = @splat(0);
17 inline for (0..D.lane_count) |index| {
18 if (index < lanes) result[index] = input[index];
19 }
20 return result;
21 }
22
23 pub fn loadNOr(
24 comptime D: type,
25 inactive: D.Vector,
26 input: []const D.Lane,
27 count: usize,
28 ) D.Vector {
29 const lanes = @min(count, D.lane_count);
30 std.debug.assert(input.len >= lanes);
31 var result = inactive;
32 inline for (0..D.lane_count) |index| {
33 if (index < lanes) result[index] = input[index];
34 }
35 return result;
36 }
37
38 pub fn storeN(
39 comptime D: type,
40 value: D.Vector,
41 output: []D.Lane,
42 count: usize,
43 ) void {
44 const lanes = @min(count, D.lane_count);
45 std.debug.assert(output.len >= lanes);
46 inline for (0..D.lane_count) |index| {
47 if (index < lanes) output[index] = value[index];
48 }
49 }
50
51 pub fn loadDup128(comptime D: type, input: []const D.Lane) D.Vector {
52 const block_lanes = @min(D.lane_count, 16 / @sizeOf(D.Lane));
53 std.debug.assert(input.len >= block_lanes);
54 var result: D.Vector = undefined;
55 inline for (0..D.lane_count) |index| {
56 result[index] = input[index % block_lanes];
57 }
58 return result;
59 }
60
61 pub fn maskedLoadOr(
62 comptime D: type,
63 inactive: D.Vector,
64 mask: D.Mask,
65 input: []const D.Lane,
66 ) D.Vector {
67 var result = inactive;
68 inline for (0..D.lane_count) |index| {
69 if (mask[index]) {
70 std.debug.assert(index < input.len);
71 result[index] = input[index];
72 }
73 }
74 return result;
75 }
76
77 pub fn maskedLoad(
78 comptime D: type,
79 mask: D.Mask,
80 input: []const D.Lane,
81 ) D.Vector {
82 return maskedLoadOr(D, @splat(0), mask, input);
83 }
84
85 pub fn insertIntoUpper(
86 comptime D: type,
87 input: []const D.Lane,
88 value: D.Vector,
89 ) D.Vector {
90 const half = D.lane_count / 2;
91 std.debug.assert(input.len >= half);
92 var result = value;
93 inline for (0..half) |index| result[index + half] = input[index];
94 return result;
95 }
96
97 pub fn safeFillN(
98 comptime D: type,
99 count: usize,
100 value: D.Lane,
101 output: []D.Lane,
102 ) void {
103 const lanes = @min(count, D.lane_count);
104 std.debug.assert(output.len >= lanes);
105 for (output[0..lanes]) |*lane_value| lane_value.* = value;
106 }
107
108 pub fn safeCopyN(
109 comptime D: type,
110 count: usize,
111 input: []const D.Lane,
112 output: []D.Lane,
113 ) void {
114 const lanes = @min(count, D.lane_count);
115 std.debug.assert(input.len >= lanes);
116 std.debug.assert(output.len >= lanes);
117 std.mem.copyForwards(D.Lane, output[0..lanes], input[0..lanes]);
118 }
119
120 pub fn truncateStore(comptime D: type, value: D.Vector, output: anytype) void {
121 const OutputLane = outputLane(@TypeOf(output));
122 if (comptime @typeInfo(D.Lane) != .int or @typeInfo(OutputLane) != .int or
123 @typeInfo(D.Lane).int.signedness != .unsigned or
124 @typeInfo(OutputLane).int.signedness != .unsigned or
125 @bitSizeOf(OutputLane) >= @bitSizeOf(D.Lane))
126 {
127 @compileError("truncateStore requires a narrower unsigned integer output");
128 }
129 std.debug.assert(output.len >= D.lane_count);
130 inline for (0..D.lane_count) |index| output[index] = @truncate(value[index]);
131 }
132
133 pub fn stream(comptime D: type, value: D.Vector, output: []D.Lane) void {
134 store(D, value, output);
135 }
136
137 pub fn blendedStore(
138 comptime D: type,
139 value: D.Vector,
140 mask: D.Mask,
141 output: []D.Lane,
142 ) void {
143 std.debug.assert(output.len >= D.lane_count);
144 inline for (0..D.lane_count) |index| {
145 if (mask[index]) output[index] = value[index];
146 }
147 }
148
149 pub fn loadInterleaved2(comptime D: type, input: []const D.Lane) [2]D.Vector {
150 std.debug.assert(input.len >= D.lane_count * 2);
151 var result: [2]D.Vector = undefined;
152 inline for (0..D.lane_count) |lane_index| {
153 inline for (0..2) |vector_index| {
154 result[vector_index][lane_index] = input[lane_index * 2 + vector_index];
155 }
156 }
157 return result;
158 }
159
160 pub fn loadInterleaved3(comptime D: type, input: []const D.Lane) [3]D.Vector {
161 std.debug.assert(input.len >= D.lane_count * 3);
162 var result: [3]D.Vector = undefined;
163 inline for (0..D.lane_count) |lane_index| {
164 inline for (0..3) |vector_index| {
165 result[vector_index][lane_index] = input[lane_index * 3 + vector_index];
166 }
167 }
168 return result;
169 }
170
171 pub fn loadInterleaved4(comptime D: type, input: []const D.Lane) [4]D.Vector {
172 std.debug.assert(input.len >= D.lane_count * 4);
173 var result: [4]D.Vector = undefined;
174 inline for (0..D.lane_count) |lane_index| {
175 inline for (0..4) |vector_index| {
176 result[vector_index][lane_index] = input[lane_index * 4 + vector_index];
177 }
178 }
179 return result;
180 }
181
182 pub fn storeInterleaved2(
183 comptime D: type,
184 a: D.Vector,
185 b: D.Vector,
186 output: []D.Lane,
187 ) void {
188 std.debug.assert(output.len >= D.lane_count * 2);
189 const values = [2]D.Vector{ a, b };
190 inline for (0..D.lane_count) |lane_index| {
191 inline for (0..2) |vector_index| {
192 output[lane_index * 2 + vector_index] = values[vector_index][lane_index];
193 }
194 }
195 }
196
197 pub fn storeInterleaved3(
198 comptime D: type,
199 a: D.Vector,
200 b: D.Vector,
201 c: D.Vector,
202 output: []D.Lane,
203 ) void {
204 std.debug.assert(output.len >= D.lane_count * 3);
205 const values = [3]D.Vector{ a, b, c };
206 inline for (0..D.lane_count) |lane_index| {
207 inline for (0..3) |vector_index| {
208 output[lane_index * 3 + vector_index] = values[vector_index][lane_index];
209 }
210 }
211 }
212
213 pub fn storeInterleaved4(
214 comptime D: type,
215 a: D.Vector,
216 b: D.Vector,
217 c: D.Vector,
218 d: D.Vector,
219 output: []D.Lane,
220 ) void {
221 std.debug.assert(output.len >= D.lane_count * 4);
222 const values = [4]D.Vector{ a, b, c, d };
223 inline for (0..D.lane_count) |lane_index| {
224 inline for (0..4) |vector_index| {
225 output[lane_index * 4 + vector_index] = values[vector_index][lane_index];
226 }
227 }
228 }
229
230 fn outputLane(comptime Output: type) type {
231 return switch (@typeInfo(Output)) {
232 .pointer => |info| switch (@typeInfo(info.child)) {
233 .array => |array_info| array_info.child,
234 else => info.child,
235 },
236 else => @compileError("truncateStore output must be a slice or pointer"),
237 };
238 }
239
240 test "full load and store preserve every lane" {
241 const simd = @import("root.zig");
242 const D = simd.FixedTag(u32, 4);
243 const input = [_]u32{ 1, 2, 3, 4 };
244 var output = [_]u32{ 0, 0, 0, 0 };
245 store(D, load(D, &input), &output);
246 try std.testing.expectEqualSlices(u32, &input, &output);
247 }
248
249 test "partial and masked memory operations leave inactive lanes alone" {
250 const simd = @import("root.zig");
251 const D = simd.FixedTag(u16, 4);
252 const input = [_]u16{ 1, 2, 3, 4 };
253 const loaded = loadNOr(D, @splat(9), &input, 2);
254 const expected: D.Vector = .{ 1, 2, 9, 9 };
255 try std.testing.expect(@reduce(.And, loaded == expected));
256
257 var output = [_]u16{ 7, 7, 7, 7 };
258 const mask: D.Mask = .{ true, false, true, false };
259 blendedStore(D, load(D, &input), mask, &output);
260 try std.testing.expectEqualSlices(u16, &.{ 1, 7, 3, 7 }, &output);
261 }
262
263 test "Highway bounded memory clamps counts and duplicates 128-bit blocks" {
264 const simd = @import("root.zig");
265 const D = simd.FixedTag(u32, 8);
266 const input = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 };
267 try std.testing.expect(@reduce(.And, loadN(D, &input, 99) == @as(D.Vector, input)));
268 try std.testing.expect(@reduce(.And, loadDup128(D, &input) ==
269 @as(D.Vector, .{ 1, 2, 3, 4, 1, 2, 3, 4 })));
270 const mask: D.Mask = .{ true, false, true, false, false, false, false, false };
271 try std.testing.expect(@reduce(.And, maskedLoadOr(D, @splat(9), mask, input[0..3]) ==
272 @as(D.Vector, .{ 1, 9, 3, 9, 9, 9, 9, 9 })));
273 try std.testing.expect(@reduce(.And, insertIntoUpper(D, input[0..4], @splat(7)) ==
274 @as(D.Vector, .{ 7, 7, 7, 7, 1, 2, 3, 4 })));
275 }
276
277 test "Highway safe copy fill and truncated store preserve bounds" {
278 const simd = @import("root.zig");
279 const D = simd.FixedTag(u16, 4);
280 const input = [_]u16{ 1, 2, 3, 4 };
281 var output = [_]u16{ 9, 9, 9, 9, 9 };
282 safeCopyN(D, 2, &input, &output);
283 try std.testing.expectEqualSlices(u16, &.{ 1, 2, 9, 9, 9 }, &output);
284 safeFillN(D, 99, 6, &output);
285 try std.testing.expectEqualSlices(u16, &.{ 6, 6, 6, 6, 9 }, &output);
286 var bytes: [4]u8 = undefined;
287 truncateStore(D, @as(D.Vector, .{ 0x123, 0x2ff, 3, 4 }), &bytes);
288 try std.testing.expectEqualSlices(u8, &.{ 0x23, 0xff, 3, 4 }, &bytes);
289 }
290
291 fn verifyStream(comptime T: type) !void {
292 const simd = @import("root.zig");
293 const D = simd.FixedTag(T, 4);
294 var expected: [4]T = undefined;
295 for (&expected, 0..) |*value, index| {
296 value.* = switch (@typeInfo(T)) {
297 .int => @intCast(index + 1),
298 .float => @floatFromInt(index + 1),
299 else => unreachable,
300 };
301 }
302 var output: [8]T = @splat(0);
303 stream(D, @as(D.Vector, expected), &output);
304 simd.flushStream();
305 try std.testing.expectEqualSlices(T, &expected, output[0..4]);
306 try std.testing.expectEqualSlices(T, &@as([4]T, @splat(0)), output[4..]);
307 }
308
309 test "Highway vector stream preserves typed lanes and rounded bounds" {
310 inline for (.{ u32, i32, u64, i64, f32, f64 }) |T| try verifyStream(T);
311 }
312
313 fn verifyInterleavedLaneType(comptime T: type) !void {
314 const simd = @import("root.zig");
315 const D = simd.FixedTag(T, 4);
316 const zero: D.Vector = @splat(0);
317 var output: [D.lane_count * 4]T = undefined;
318 storeInterleaved4(D, zero, zero, zero, zero, &output);
319 const loaded = loadInterleaved4(D, &output);
320 inline for (loaded) |value| try std.testing.expect(@reduce(.And, value == zero));
321 }
322
323 test "Highway interleaved memory instantiates every lane type" {
324 inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {
325 try verifyInterleavedLaneType(T);
326 }
327 }
328
329 test "Highway two three and four channel memory uses lane-major order" {
330 const simd = @import("root.zig");
331 const D = simd.FixedTag(u16, 4);
332 const a: D.Vector = .{ 10, 11, 12, 13 };
333 const b: D.Vector = .{ 20, 21, 22, 23 };
334 const c: D.Vector = .{ 30, 31, 32, 33 };
335 const d: D.Vector = .{ 40, 41, 42, 43 };
336
337 var two: [8]u16 = undefined;
338 storeInterleaved2(D, a, b, &two);
339 try std.testing.expectEqualSlices(u16, &.{ 10, 20, 11, 21, 12, 22, 13, 23 }, &two);
340 const loaded_two = loadInterleaved2(D, &two);
341 try std.testing.expect(@reduce(.And, loaded_two[0] == a));
342 try std.testing.expect(@reduce(.And, loaded_two[1] == b));
343
344 var three: [12]u16 = undefined;
345 storeInterleaved3(D, a, b, c, &three);
346 try std.testing.expectEqualSlices(u16, &.{
347 10, 20, 30, 11, 21, 31, 12, 22, 32, 13, 23, 33,
348 }, &three);
349 const loaded_three = loadInterleaved3(D, &three);
350 try std.testing.expect(@reduce(.And, loaded_three[2] == c));
351
352 var four: [16]u16 = undefined;
353 storeInterleaved4(D, a, b, c, d, &four);
354 try std.testing.expectEqualSlices(u16, &.{
355 10, 20, 30, 40, 11, 21, 31, 41,
356 12, 22, 32, 42, 13, 23, 33, 43,
357 }, &four);
358 const loaded_four = loadInterleaved4(D, &four);
359 try std.testing.expect(@reduce(.And, loaded_four[3] == d));
360 }