Missing SIMD shift instructions by non-scalar amounts

[Vanilagy]

The current SIMD integer shift instructions (and even the proposed ones in flexible-vectors) only ever do "shift by scalar". However, CPUs do have instructions for "shift a vector by another vector", like VQSHL as an example.

I was trying to do parallel integer processing, but since it's extremely shift-heavy and the shifts are per-element variable, it's basically equivalent to the scalar code performance-wise (or worse, actually).

My question is: why do these instructions not exist? Especially in WASM where, if the underlying hardware doesn't have a corresponding instruction, the runtime can just trivially emit a scalarized version?

[ppenzin]

I don't think the use case of shifting lanes of one vector by values of corresponding lanes in another vector came up in prior SIMD discussions. Can you share a bit more about the algorithm you are porting?

[Vanilagy]

@ppenzin I'm working on a WASM ProRes decoder and at the core of it is this function (here written in Zig) that parses an entropy-coded code word:

inline fn parseCode(word: u64, params: u64) ParsedCode {
    const mp: u64 = params & 0b11;
    const g: u64 = (params >> 2) & 0b111;
    const r: u64 = params >> 5;

    const n: u64 = @clz(word);
    const is_big = n > mp;

    const base = @as(u64, @min(n, mp + 1)) << @as(u6, @intCast(r));

    const bits_big = (n << 1) +% g -% mp;
    const bits_small = n + 1 + r;
    const bits = if (is_big) bits_big else bits_small;
    const sub = @as(u64, 1) << (if (is_big) @intCast(g) else @intCast(r));
    const raw = word >> @as(u6, @intCast(64 - bits));

    const result = base +% raw -% sub;

    return .{
        .value = result,
        .bits = bits,
    };
}

I was vectorizing this parser so that it parses four streams at the same time - possible because the above function is branchless (although it does contain some cmovs). The only non-vectorizable instruction here is clz, but that's fine with me, I can dip out of "vector land" once. But, the rest of the code contains a bunch of shifts by dynamic amounts for each lane because of this. Which unfortunately WASM can't model.

[ppenzin]

You are vectorizing this manually, right?

For 128 bit SIMD 64 bit chunks might not yield very impressive speed ups, but that I what I guess we have flexible vector proposal for.

What instructions are you getting on native targets, particularly x86?

[Vanilagy]

What do you mean by "manually"? I am writing Zig code that uses @Vector, which LLVM compiles to WASM SIMD instructions. I use 64 bit for the scalar version, but the code technically only needs 32, so I'd use a i32x4 in my use case, or even two i32x4 at the same time. (In Zig I can just specify a @Vector(8, u32) and LLVM deals with it properly, splitting it up). But yes, flexible vectors would of course make this even better.

I have no compiled it for native targets so I can't say.

[ppenzin]

I meant something like intrinsics. Curious how this works in Zig, for example in C-like languages build vector from array accesses.

The reason I am asking about native platforms is because we had situations where one platform has some weird corner case preventing instruction lowering.

[Vanilagy]

Zig has one of these "higher-level" SIMD implementations where you completely avoid intrinsics and operate on abstract "Vectors" that the backend then compiles to the correct instructions. Vectors can have arbitrary length (well, up to like 2^32-1).

So, a lane-wise shift is roughly written as:

const a: @Vector(4, u32) = ...;
const b: @Vector(4, u32) = ...;
const result = a << b;

Shift by scalar would be:

const a: @Vector(4, u32) = ...;
const result = a << @splat(7);