zig
zig copied to clipboard
Indexing arrays with vectors (gather)
This is a proposal to add a function that enables you to index an array with a vector, producing a vector of the array's elements as its output.
Here is a workaround in the meantime:
const std = @import("std");
const expect = std.testing.expect;
fn gather(slice: anytype, index: anytype) @Vector(
@typeInfo(@TypeOf(index)).Vector.len,
@typeInfo(@TypeOf(slice)).Pointer.child,
) {
const vector_len = @typeInfo(@TypeOf(index)).Vector.len;
const Elem = @typeInfo(@TypeOf(slice)).Pointer.child;
var result: [vector_len]Elem = undefined;
comptime var vec_i = 0;
inline while (vec_i < vector_len) : (vec_i += 1) {
result[vec_i] = slice[index[vec_i]];
}
return result;
}
test "vector gather" {
const hello: []const u8 = "hello world";
const index: @Vector(3, usize) = .{ 1, 3, 5 };
const result = gather(hello, index);
try expect(@TypeOf(result) == @Vector(3, u8));
try expect(result[0] == 'e');
try expect(result[1] == 'l');
try expect(result[2] == ' ');
}
This sounds more like a "shuffle" op, technically
An array-vector shuffle can be implemented in terms of an array-vector address-of and a vector gather
For example, the way that LLVM supports this is by supporting &array[vec]
(yields a vector of pointers) and a gather/scatter that takes a vector of pointers to load to/from (with a mask)
Do we want to expose those primitive ops (array-vector indexing, vector-vector masked gather) directly? Or is an array-vector shuffle enough/preferable?