Sharing code between scalar and vector types

Help
1

Couple questions about how folk are approaching working with vector types for SIMD work.

Say I have some complex function that takes a number of f64 arguments. I want to avoid writing two implementations of the same logic, one for scalar f64 and one for @Vector(n, f64). I have a few questions / thoughts on this below.

One solution: use a signature that takes anytype, and use type reflection to detect vector types. This is doable, but stuff like scalar mutliplication is scalar * x in scalar land and the lovely @as(Vec, @splat(3.0)) * x in vector-land, so you have to do a lot of these comptime checks.

Another solution I’ve been playing with is writing everything as vector only, and special casing scalar calls as vectors of length 1. For example

fn typedMultiply(T: type, x: T, y: T) T {
    return x * y;
}


export fn multiplyAsVec(x: f64, y: f64) f64 {
    return typedMultiply(@Vector(1, f64), @splat(x), @splat(y))[0];
    
}

export fn multiplyAsScalar(x: f64, y: f64) f64 {
    return typedMultiply(f64, x, y);
}

I was interested in whether there is overhead on doing the vector to scalar conversion - so ran this through to look at the assembly. Interestingly, it’s really exactly the same (to my very, very untrained eye).

example.multiplyAsScalar:
        push    rbp
        mov     rbp, rsp
        sub     rsp, 304
        vmovsd  qword ptr [rbp - 296], xmm0
        vmovsd  qword ptr [rbp - 288], xmm1
        lea     rax, [rbp - 280]
        mov     qword ptr [rbp - 16], rax
        mov     qword ptr [rbp - 8], 32
        mov     qword ptr [rbp - 24], 0
        lea     rdi, [rbp - 24]
        call    example.typedMultiply__anon_482
        add     rsp, 304
        pop     rbp
        ret

example.typedMultiply__anon_482:
        push    rbp
        mov     rbp, rsp
        sub     rsp, 16
        vmovsd  qword ptr [rbp - 16], xmm0
        vmovsd  qword ptr [rbp - 8], xmm1
        vmulsd  xmm0, xmm0, xmm1
        add     rsp, 16
        pop     rbp
        ret

example.multiplyAsVec:
        push    rbp
        mov     rbp, rsp
        sub     rsp, 304
        vmovsd  qword ptr [rbp - 304], xmm0
        vmovsd  qword ptr [rbp - 296], xmm1
        lea     rax, [rbp - 288]
        mov     qword ptr [rbp - 24], rax
        mov     qword ptr [rbp - 16], 32
        mov     qword ptr [rbp - 32], 0
        lea     rdi, [rbp - 32]
        call    example.typedMultiply__anon_491
        vmovsd  qword ptr [rbp - 8], xmm0
        vmovsd  xmm0, qword ptr [rbp - 8]
        add     rsp, 304
        pop     rbp
        ret

example.typedMultiply__anon_491:
        push    rbp
        mov     rbp, rsp
        sub     rsp, 16
        vmovsd  qword ptr [rbp - 16], xmm0
        vmovsd  qword ptr [rbp - 8], xmm1
        vmulsd  xmm0, xmm0, xmm1
        add     rsp, 16
        pop     rbp
        ret

My question then is: is this a sensible pattern? Or am I going to run into gotchas where vector-based code will use slower instruction sets on scalars or things like this?

Please do throw in any other tips for this kind of work as well :slight_smile: thanks all!

2 Likes
2

Your assembly outplut looks like it is compiled in debug mode.
You are right that your two functions compile to the same machine code. This gets even more obvious when compiling with ReleaseSafe:

example.multiplyAsScalar:
        push    rbp
        mov     rbp, rsp
        vmulsd  xmm0, xmm0, xmm1
        pop     rbp
        ret

multiplyAsScalar = example.multiplyAsScalar
multiplyAsVec = example.multiplyAsScalar

This is because SIMD vectors in zig are not really special. They do not force the compiler to use SIMD registers. They just keep data in a shape for which the optimizer probably is able to use simd instructions if it sees fit.

Under the right cicumstances scalar zig code can also generate simd instructions. The following will most likely result in vectorized machine code if f is just some scalar math.

fn apply2(f: fn (f64, f64) f64, x: @Vector(8, f64), y: @Vector(8, f64)) @Vector(8, f64) {
    var res: [8]f64 = undefined;
    const x_arr: [8]f64 = x;
    const y_arr: [8]f64 = y;
    for (&res, x_arr, y_arr) |*r, xe, ye| r.* = f(xe, ye);
    return res;
}

Here is a version of apply that works for functions that take arbitrarily many args:

Note that this produces exactly the same assembly if I replace all uses of @Vector(n, T) with [n]T.

4 Likes