@shuffle lookup table without branching

[vale valley]

So I have a comptime lookup table of 256 shuffle masks. At runtime, I calculate a u8 to index into this table and select the correct mask. The only way I've found to do this is with an inline switch:

const result = switch (code) {
  inline else => |c| @shuffle(..., ..., ..., masks[c]),
};

The problem is, this generates a jump table which I'm guessing makes the branch predictor very mad indeed...

Is there any way to work around this without hand-rolling inline asm intrinsics for each instruction set?

[bleak sleet]

do you mean you get decreased performance when you do this? you could try writing a runtime version of shuffle, since the mask has to be known at comptime

[vale valley]

yes, greatly decreased performance vs a lookup table and single pshufb (or equivalent) instruction.

I could write a runtime one, but it would have to switch on architecture and use inline asm, which I was hoping to avoid.

[bleak sleet]

maybe try writing the shuffle function in regular zig? it might not be as fast as asm but would most likely be faster than your current solution

[bleak sleet]

this is the best i can do for no branching 😅:

fn ShuffleReturn(comptime T: type) type {
    return @Vector(@typeInfo(T).Vector.len, @typeInfo(T).Vector.child);
}

fn shuffle(a: anytype, b: @TypeOf(a), mask: @Vector(@typeInfo(@TypeOf(a)).Vector.len, i32)) ShuffleReturn(@TypeOf(a)) {
    const Child = @typeInfo(@TypeOf(a)).Vector.child;
    const len = @typeInfo(@TypeOf(a)).Vector.len;

    var vals: [len * 2]Child = undefined;
    @memcpy(vals[0..len], &@as([len]Child, b));
    @memcpy(vals[len..], &@as([len]Child, a));
    std.mem.reverse(Child, vals[0..len]);
    var out: [len]Child = undefined;

    for (@as([len]i32, mask), &out) |m, *o| {
        o.* = vals[@intCast(m + len)];
    }

    return out;
}

ah wait the original shuffle allows differing lengths, and some undefined stuff

heres one that allows different size vectors:

fn shuffle(comptime E: type, a: anytype, b: anytype, mask: anytype) @Vector(mask.len, E) {
    const a_len = @typeInfo(@TypeOf(a)).Vector.len;
    const b_len = @typeInfo(@TypeOf(b)).Vector.len;

    var vals: [a_len + b_len]E = undefined;
    @memcpy(vals[0..b_len], &@as([b_len]E, b));
    @memcpy(vals[b_len..], &@as([a_len]E, a));
    std.mem.reverse(E, vals[0..b_len]);

    var out: [mask.len]E = undefined;

    for (@as([mask.len]i32, mask), &out) |m, *o| {
        o.* = vals[@intCast(m + b_len)];
    }

    return out;
}

[empty raptor]

if you write it as a for loop, can LLVM figure it out and turn it into a single shuffle?

[fickle mortar]

Can't you just to const result = @shuffle(..., ..., ..., masks[code])? Without more context I don't see why this wouldn't work.

[vale valley]

Because "code" is only runtime known. So although masks is comptime known, masks[code] is no longer comptime.

I think I may have figured out how to convince LLVM to optimise this into a shuffle operation. I believe the key is that the mask has to be a full permutation. e.g. 1423 would work, but 1122 would not. Still fiddling around atm

[fickle mortar]

I don't see why you need masks[code] to be comptime. Having masks[code] be runtime shouldn't introduce a branch, it should just lower into a single pshufb or similar right? I feel like I must be misunderstanding something...

[vale valley]

Ah, so tThe mask parameter of @shuffle is required to be comptime

[hot rain]

Maybe you should open an issue for that.

[empty raptor]

I think ASM frankly seems like a good idea here.

How many targets do you support anyway, and do all of them really need to be super fast?
Just make it optimized on the mabye two or three targets you actually care about, probably x86_64 and ARM, and screw the rest and let them use a jump table.

[hot rain]

IMO, using inline assembly is rarely a good idea. Zig should be able to generate reasonably quick assembly for basic operations with its vector primitives. When it can't, that's a language problem.

[quartz herald]

Hello, you basically need to use asm for this. Luckily there are only a few flavors of pshufb in the world: 16-wide, 32-wide, and 64-wide on x64, 16-wide tbl on NEON (the wider tbl's are something different), and some weird thing in upcoming SVE.

[quartz herald]

LLVM's vector stuff, which Zig's vector stuff is based on at least for historical reasons, does not support dynamic shuffles at all

[hot rain]

💀

[fickle mortar]

Hmm, interesting - this seems like something Zig should want to fix given that the actual instructions these days (at least on x86_64) don't have that restriction

[vale valley]

Yup, I ended up going the ASM route.

For my limited use-case, this seems to work. There are probably corner cases where the behaviour differs.

// For all vector widths; Arm anyway zeroes if >= 0x10.
inline fn Aarch64_Neon(bytes: @Vector(16, u8), indices: @Vector(16, i8)) @Vector(16, u8) {
    return asm ("tbl.16b %[ret], { %[v0] }, %[v1]"
        : [ret] "=w" (-> @Vector(16, u8)),
        : [v0] "w" (bytes),
          [v1] "w" (indices),
    );
}

inline fn X64_SSE3(bytes: @Vector(16, u8), indices: @Vector(16, i8)) @Vector(16, u8) {
    var result = bytes;
    asm volatile ("pshufb %[indices], %[bytes]"
        : [bytes] "+x" (result),
        : [indices] "x" (indices),
    );
    return result;
}

[vale valley]

In case anyone is interested, I just open-sourced the project that was using this: https://github.com/fulcrum-so/streamvbyte-zig

Conclusion is, Zig would be incredible for building these types of cross-platform compression algorithms if SIMD was a little more powerful.

Perhaps there's scope for a Highway-like Zig library to experiment with providing reasonable cross-platform semantics without the high bar necessary for inclusion in the stdlib.