index array with a vector

[native mist]

Hi! :D

I have a vector of indexes into an array and I would like to turn that vector of indexes into a vector of the associated values from some array like this:

const std = @import("std");
const assert = std.debug.assert;

pub fn main() !void {
    var arr = [_]i32{ 1, 2, 3, 4, 5, 6, 7 };
    var a = @Vector(4, i32){ 6, 2, 0, 4 };

    var b: @Vector(4, i32) = arr[a];

    assert(b == @Vector(4, i32){ 7, 3, 1, 5 });
}

Is is possible? would there be a benefit from a doing this instead of some non-vectorized manual load?

The usecase is getting piece types from a the board in my chess engine btw: I have ~16 to ~64 square indexes that I need to turn into piece types (enum) before doing move ordering and I thought that loading the piece types in bulk like this would be faster than just some loop.

Thanks for any help :D

[dire gazelle]

That's not implemented yet, but it is an accepted proposal(but I guess it will probably be a builtin @gather(arr, vec) rather than arr[vec]): https://github.com/ziglang/zig/issues/12815
For now you have to do it manually (you can the workaround function andrew posted in that issue)

[rain smelt]

note that you can use @shuffle() if the mask is comptime known:

const std = @import("std");
const t = std.testing;

test {
    // if the mask is comptime known you can use @shuffle()
    var arr = [_]i32{ 1, 2, 3, 4, 5, 6, 7 };
    // this must be const, not var
    const a = @Vector(4, i32){ 6, 2, 0, 4 }; 
    const b = @shuffle(i32, arr, undefined, a);
    try t.expectEqualSlices(i32, &.{ 7, 3, 1, 5 }, &@as([4]i32, b));
}

otherwise, if the mask isn't comptime known like your example, you'll have use intrinsics or assembly at the moment. here is an example using vpshufb in simdjzon https://github.com/travisstaloch/simdjzon/blob/main/src/c_intrinsics.zig#L64

[rain smelt]

i'd encourage you to play around with godbolt and see if you can convince it to vectorize your loops. here's a starting point: https://godbolt.org/z/ndj5aoWox

also @pure pagoda might have some good advice

[pure pagoda]

Interesting, I wouldn't have thought of using @shuffle here because I would have thought of the array operand as "an array"