zig icon indicating copy to clipboard operation
zig copied to clipboard

Indexing arrays with vectors (gather)

Open DiskPoppy opened this issue 2 years ago • 2 comments

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.

DiskPoppy avatar Sep 11 '22 21:09 DiskPoppy

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] == ' ');
}

andrewrk avatar Sep 11 '22 21:09 andrewrk

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?

topolarity avatar Sep 12 '22 16:09 topolarity