lib/simd/src/indexed.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn scatterIndex(
4 comptime D: type,
5 value: D.Vector,
6 base: []D.Lane,
7 indices: anytype,
8 ) void {
9 validateIndices(D, @TypeOf(indices));
10 inline for (0..D.lane_count) |lane_index| {
11 const destination = checkedIndex(indices[lane_index]);
12 std.debug.assert(destination < base.len);
13 base[destination] = value[lane_index];
14 }
15 }
16
17 pub fn scatterIndexN(
18 comptime D: type,
19 value: D.Vector,
20 base: []D.Lane,
21 indices: anytype,
22 count: usize,
23 ) void {
24 validateIndices(D, @TypeOf(indices));
25 const lanes = @min(count, D.lane_count);
26 inline for (0..D.lane_count) |lane_index| {
27 if (lane_index < lanes) {
28 const destination = checkedIndex(indices[lane_index]);
29 std.debug.assert(destination < base.len);
30 base[destination] = value[lane_index];
31 }
32 }
33 }
34
35 pub fn maskedScatterIndex(
36 comptime D: type,
37 value: D.Vector,
38 mask: D.Mask,
39 base: []D.Lane,
40 indices: anytype,
41 ) void {
42 validateIndices(D, @TypeOf(indices));
43 inline for (0..D.lane_count) |lane_index| {
44 if (mask[lane_index]) {
45 const destination = checkedIndex(indices[lane_index]);
46 std.debug.assert(destination < base.len);
47 base[destination] = value[lane_index];
48 }
49 }
50 }
51
52 pub fn scatterOffset(
53 comptime D: type,
54 value: D.Vector,
55 base: []D.Lane,
56 offsets: anytype,
57 ) void {
58 validateIndices(D, @TypeOf(offsets));
59 inline for (0..D.lane_count) |lane_index| {
60 const destination = offsetIndex(D.Lane, offsets[lane_index]);
61 std.debug.assert(destination < base.len);
62 base[destination] = value[lane_index];
63 }
64 }
65
66 pub fn gatherIndex(comptime D: type, base: []const D.Lane, indices: anytype) D.Vector {
67 validateIndices(D, @TypeOf(indices));
68 var result: D.Vector = undefined;
69 inline for (0..D.lane_count) |lane_index| {
70 const source = checkedIndex(indices[lane_index]);
71 std.debug.assert(source < base.len);
72 result[lane_index] = base[source];
73 }
74 return result;
75 }
76
77 pub fn gatherIndexN(
78 comptime D: type,
79 base: []const D.Lane,
80 indices: anytype,
81 count: usize,
82 ) D.Vector {
83 return gatherIndexNOr(D, @splat(0), base, indices, count);
84 }
85
86 pub fn gatherIndexNOr(
87 comptime D: type,
88 inactive: D.Vector,
89 base: []const D.Lane,
90 indices: anytype,
91 count: usize,
92 ) D.Vector {
93 const lanes = @min(count, D.lane_count);
94 var mask: D.Mask = @splat(false);
95 inline for (0..D.lane_count) |index| mask[index] = index < lanes;
96 return maskedGatherIndexOr(D, inactive, mask, base, indices);
97 }
98
99 pub fn maskedGatherIndexOr(
100 comptime D: type,
101 inactive: D.Vector,
102 mask: D.Mask,
103 base: []const D.Lane,
104 indices: anytype,
105 ) D.Vector {
106 validateIndices(D, @TypeOf(indices));
107 var result = inactive;
108 inline for (0..D.lane_count) |lane_index| {
109 if (mask[lane_index]) {
110 const source = checkedIndex(indices[lane_index]);
111 std.debug.assert(source < base.len);
112 result[lane_index] = base[source];
113 }
114 }
115 return result;
116 }
117
118 pub fn maskedGatherIndex(
119 comptime D: type,
120 mask: D.Mask,
121 base: []const D.Lane,
122 indices: anytype,
123 ) D.Vector {
124 return maskedGatherIndexOr(D, @splat(0), mask, base, indices);
125 }
126
127 pub fn gatherOffset(comptime D: type, base: []const D.Lane, offsets: anytype) D.Vector {
128 validateIndices(D, @TypeOf(offsets));
129 var result: D.Vector = undefined;
130 inline for (0..D.lane_count) |lane_index| {
131 const source = offsetIndex(D.Lane, offsets[lane_index]);
132 std.debug.assert(source < base.len);
133 result[lane_index] = base[source];
134 }
135 return result;
136 }
137
138 fn checkedIndex(value: anytype) usize {
139 const T = @TypeOf(value);
140 if (comptime @typeInfo(T) != .int) @compileError("indices must be integers");
141 if (comptime @typeInfo(T).int.signedness == .signed) std.debug.assert(value >= 0);
142 return @intCast(value);
143 }
144
145 fn offsetIndex(comptime T: type, offset: anytype) usize {
146 const bytes = checkedIndex(offset);
147 std.debug.assert(bytes % @sizeOf(T) == 0);
148 return bytes / @sizeOf(T);
149 }
150
151 fn validateIndices(comptime D: type, comptime V: type) void {
152 if (comptime @typeInfo(V) != .vector or @typeInfo(V).vector.len != D.lane_count or
153 @typeInfo(@typeInfo(V).vector.child) != .int)
154 {
155 @compileError("indices must be an integer vector with one index per lane");
156 }
157 }
158
159 test "Highway indexed gather and scatter use lane order and clamped counts" {
160 const simd = @import("root.zig");
161 const D = simd.FixedTag(u32, 4);
162 const I = simd.FixedTag(i32, 4);
163 const indices: I.Vector = .{ 3, 0, 5, 2 };
164 const source = [_]u32{ 10, 11, 12, 13, 14, 15 };
165 try std.testing.expect(@reduce(.And, gatherIndex(D, &source, indices) ==
166 @as(D.Vector, .{ 13, 10, 15, 12 })));
167 try std.testing.expect(@reduce(.And, gatherIndexN(D, &source, indices, 2) ==
168 @as(D.Vector, .{ 13, 10, 0, 0 })));
169 var output = [_]u32{ 0, 0, 0, 0, 0, 0 };
170 scatterIndexN(D, .{ 1, 2, 3, 4 }, &output, indices, 3);
171 try std.testing.expectEqualSlices(u32, &.{ 2, 0, 0, 1, 0, 3 }, &output);
172 }
173
174 test "Highway masked and byte-offset memory avoids inactive indices" {
175 const simd = @import("root.zig");
176 const D = simd.FixedTag(f64, 2);
177 const I = simd.FixedTag(i64, 2);
178 const source = [_]f64{ 1, 2, 3, 4 };
179 const mask: D.Mask = .{ true, false };
180 const indices: I.Vector = .{ 2, -1 };
181 try std.testing.expect(@reduce(.And, maskedGatherIndexOr(D, @splat(9), mask, &source, indices) ==
182 @as(D.Vector, .{ 3, 9 })));
183 const offsets: I.Vector = .{ 24, 0 };
184 try std.testing.expect(@reduce(.And, gatherOffset(D, &source, offsets) ==
185 @as(D.Vector, .{ 4, 1 })));
186 var output = [_]f64{ 0, 0, 0, 0 };
187 scatterOffset(D, .{ 7, 8 }, &output, offsets);
188 try std.testing.expectEqualSlices(f64, &.{ 8, 0, 0, 7 }, &output);
189 }