Efficient 3x3 matrix multiplication using SIMD

Help
1

Hi all!

I am making a simple matrix math implementation for my own game. I would like to use SIMD and the vector type for these operations. Right now, my 4x4 matrix multiplication is as following:

pub const Mat4x4 = [4][4]f32;

fn matMul4(m1: Mat4x4, m2: Mat4x4) Mat4x4 {
    var m3: Mat4x4 = undefined;
    const Vec = @Vector(4, f32);
    const m1v: [4]Vec = .{
        @as(Vec, m1[0]),
        @as(Vec, m1[1]),
        @as(Vec, m1[2]),
        @as(Vec, m1[3]),
    };

    inline for (0..4) |i| {
        const s0: Vec = @splat(m2[i][0]);
        const s1: Vec = @splat(m2[i][1]);
        const s2: Vec = @splat(m2[i][2]);
        const s3: Vec = @splat(m2[i][3]);

        m3[i] = s0 * m1v[0] + s1 * m1v[1] + s2 * m1v[2] + s3 * m1v[3];
    }

    return m3;
}

First question, in general, would this be a good way to implement this operation?

Now, I would like to make the 3x3 version. My first instinct is to make something like that:

pub const Mat3x3 = [3][3]f32;

fn matMul3(m1: Mat3x3, m2: Mat3x3) Mat3x3 {
    var m3: Mat3x3 = undefined;
    const Vec = @Vector(3, f32);
    const m1v: [3]Vec = .{
        @as(Vec, m1[0]),
        @as(Vec, m1[1]),
        @as(Vec, m1[2]),
    };

    inline for (0..3) |i| {
        const s0: Vec = @splat(m2[i][0]);
        const s1: Vec = @splat(m2[i][1]);
        const s2: Vec = @splat(m2[i][2]);

        m3[i] = s0 * m1v[0] + s1 * m1v[1] + s2 * m1v[2];
    }

    return m3;
}

My question is: I know that in general @Vector(4,f32) is preferred since it aligns well with the SIMD operations of most machines. Would this implementation using @Vector(3,f32) be optimized correctly?

From a first research, in the docs it says:

Operations on vectors shorter than the target machine’s native SIMD size will typically compile to single SIMD instructions, (…)

So I guess it should be fine? Is there any other better way to do this?

2

Godbolt is your friend for this kind of questions

3

For sse/neon it is faster to pad your matrices to a power of two, at least in my own game that was the case, also you can do a transpose and a dot product, obviously check on godbolt, but this is probably better.
godbolt

basically the padded version does this on release fast

pub const Mat3Padded = struct {
    rows: [3]@Vector(4, f32),
};

export fn padded(
    m1: *const Mat3Padded,
    m2: *const Mat3Padded,
    out: *Mat3Padded,
) void {
    const Vec = @Vector(4, f32);

    inline for (0..3) |i| {
        const s0: Vec = @splat(m1.rows[i][0]);
        const s1: Vec = @splat(m1.rows[i][1]);
        const s2: Vec = @splat(m1.rows[i][2]);

        out.rows[i] =
            s0 * m2.rows[0] +
            s1 * m2.rows[1] +
            s2 * m2.rows[2];
    }
}

compiles down to this

example.padded:
        push    rbp
        mov     rbp, rsp
        vbroadcastss    xmm0, dword ptr [rdi]
        vbroadcastss    xmm1, dword ptr [rdi + 4]
        vbroadcastss    xmm2, dword ptr [rdi + 8]
        vmulps  xmm0, xmm0, xmmword ptr [rsi]
        vmulps  xmm1, xmm1, xmmword ptr [rsi + 16]
        vaddps  xmm0, xmm0, xmm1
        vmulps  xmm1, xmm2, xmmword ptr [rsi + 32]
        vaddps  xmm0, xmm0, xmm1
        vmovaps xmmword ptr [rdx], xmm0
        vbroadcastss    xmm0, dword ptr [rdi + 16]
        vbroadcastss    xmm1, dword ptr [rdi + 20]
        vmulps  xmm0, xmm0, xmmword ptr [rsi]
        vmulps  xmm1, xmm1, xmmword ptr [rsi + 16]
        vbroadcastss    xmm2, dword ptr [rdi + 24]
        vaddps  xmm0, xmm0, xmm1
        vmulps  xmm1, xmm2, xmmword ptr [rsi + 32]
        vaddps  xmm0, xmm0, xmm1
        vmovaps xmmword ptr [rdx + 16], xmm0
        vbroadcastss    xmm0, dword ptr [rdi + 32]
        vbroadcastss    xmm1, dword ptr [rdi + 36]
        vbroadcastss    xmm2, dword ptr [rdi + 40]
        vmulps  xmm0, xmm0, xmmword ptr [rsi]
        vmulps  xmm1, xmm1, xmmword ptr [rsi + 16]
        vmulps  xmm2, xmm2, xmmword ptr [rsi + 32]
        vaddps  xmm0, xmm0, xmm1
        vaddps  xmm0, xmm0, xmm2
        vmovaps xmmword ptr [rdx + 32], xmm0
        pop     rbp
        ret

And your version unpadded does

pub const Mat3x3 = [3][3]f32;



export fn baseline(m1: *const Mat3x3, m2: *const Mat3x3, out: *Mat3x3) void {
    const Vec = @Vector(3, f32);

    const m2v: [3]Vec = .{
        @as(Vec, m2[0]),
        @as(Vec, m2[1]),
        @as(Vec, m2[2]),
    };

    inline for (0..3) |i| {
        const s0: Vec = @splat(m1[i][0]);
        const s1: Vec = @splat(m1[i][1]);
        const s2: Vec = @splat(m1[i][2]);

        out[i] =
            s0 * m2v[0] +
            s1 * m2v[1] +
            s2 * m2v[2];
    }
}

which compiles down to :

example.baseline:
        push    rbp
        mov     rbp, rsp
        vmovsd  xmm0, qword ptr [rsi + 4]
        vmovss  xmm1, dword ptr [rsi]
        vmovlhps        xmm1, xmm1, xmm0
        vshufps xmm0, xmm1, xmm0, 216
        vmovsd  xmm1, qword ptr [rsi + 16]
        vmovss  xmm2, dword ptr [rsi + 12]
        vmovlhps        xmm2, xmm2, xmm1
        vshufps xmm1, xmm2, xmm1, 216
        vmovsd  xmm2, qword ptr [rsi + 28]
        vmovss  xmm3, dword ptr [rsi + 24]
        vmovlhps        xmm3, xmm3, xmm2
        vshufps xmm2, xmm3, xmm2, 216
        vbroadcastss    xmm3, dword ptr [rdi]
        vbroadcastss    xmm4, dword ptr [rdi + 4]
        vbroadcastss    xmm5, dword ptr [rdi + 8]
        vmulps  xmm3, xmm0, xmm3
        vmulps  xmm4, xmm1, xmm4
        vaddps  xmm3, xmm3, xmm4
        vmulps  xmm4, xmm2, xmm5
        vaddps  xmm3, xmm3, xmm4
        vmovlps qword ptr [rdx], xmm3
        vextractps      dword ptr [rdx + 8], xmm3, 2
        vbroadcastss    xmm3, dword ptr [rdi + 12]
        vbroadcastss    xmm4, dword ptr [rdi + 16]
        vbroadcastss    xmm5, dword ptr [rdi + 20]
        vmulps  xmm3, xmm0, xmm3
        vmulps  xmm4, xmm1, xmm4
        vaddps  xmm3, xmm3, xmm4
        vmulps  xmm4, xmm2, xmm5
        vaddps  xmm3, xmm3, xmm4
        vmovlps qword ptr [rdx + 12], xmm3
        vextractps      dword ptr [rdx + 20], xmm3, 2
        vbroadcastss    xmm3, dword ptr [rdi + 24]
        vbroadcastss    xmm4, dword ptr [rdi + 28]
        vbroadcastss    xmm5, dword ptr [rdi + 32]
        vmulps  xmm0, xmm0, xmm3
        vmulps  xmm1, xmm1, xmm4
        vaddps  xmm0, xmm0, xmm1
        vmulps  xmm1, xmm2, xmm5
        vaddps  xmm0, xmm0, xmm1
        vmovlps qword ptr [rdx + 24], xmm0
        vextractps      dword ptr [rdx + 32], xmm0, 2
        pop     rbp
        ret

So basically the padded version avoids all the unpacking whereas the normal version has to build the xmm registers three time

In fact you can even go further with FMA

pub const Vec4 = @Vector(4, f32);

export fn paddedFma(
    a: *const Mat3Padded,
    b: *const Mat3Padded,
    out: *Mat3Padded,
) void {
    const b0 = b.rows[0];
    const b1 = b.rows[1];
    const b2 = b.rows[2];

    inline for (0..3) |i| {
        const a0: Vec4 = @splat(a.rows[i][0]);
        const a1: Vec4 = @splat(a.rows[i][1]);
        const a2: Vec4 = @splat(a.rows[i][2]);

        var r = a0 * b0;
        r = @mulAdd(Vec4, a1, b1, r);
        r = @mulAdd(Vec4, a2, b2, r);

        out.rows[i] = r;
    }
}

example.paddedFma:
        push    rbp
        mov     rbp, rsp
        vmovaps xmm0, xmmword ptr [rsi]
        vmovaps xmm1, xmmword ptr [rsi + 16]
        vmovaps xmm2, xmmword ptr [rsi + 32]
        vmulps  xmm3, xmm0, dword ptr [rdi]{1to4}
        vfmadd231ps     xmm3, xmm1, dword ptr [rdi + 4]{1to4}
        vfmadd231ps     xmm3, xmm2, dword ptr [rdi + 8]{1to4}
        vmovaps xmmword ptr [rdx], xmm3
        vmulps  xmm3, xmm0, dword ptr [rdi + 16]{1to4}
        vfmadd231ps     xmm3, xmm1, dword ptr [rdi + 20]{1to4}
        vfmadd231ps     xmm3, xmm2, dword ptr [rdi + 24]{1to4}
        vmovaps xmmword ptr [rdx + 16], xmm3
        vmulps  xmm0, xmm0, dword ptr [rdi + 32]{1to4}
        vfmadd231ps     xmm0, xmm1, dword ptr [rdi + 36]{1to4}
        vfmadd231ps     xmm0, xmm2, dword ptr [rdi + 40]{1to4}
        vmovaps xmmword ptr [rdx + 32], xmm0
        pop     rbp
        ret

Hope i didn’t make mistakes but this is what I use found it in the valve’s Source code and absolutely stole it :slight_smile:

4

Thanks @pierrelgol, I think this pretty much answers my question. But my main problem with this approach is whether to carry these padded matrices throughout the codebase. Obviously you cannot use it for GPU-related applications, to upload to a shader for example. So in your case would you just convert it back to a [3][3]f32 when you need to actually use it?

I quickly checked your example in godbolt and converting back and forth to [3][3]f32 in each call of the paddedFma function is actually worse than just using directly the @Vector(3, f32) that I suggested earlier.

5

yeah I personally use the paddedFma version, and it works great for me, all my vectors are padded to multplie of power of 2, as for your observation, I’m curious as to why ? at least on my machine it was faster maybe it’s related to the target you are one ? Pretty sure I can find the micro benchmark somewhere in my files

EDIT: found it

src ) zig run -O ReleaseFast benchmark.zig
31 samples x 10000000 multiplications (ReleaseFast)

baseline      1.618 ns/op    1.00x baseline
padded        1.245 ns/op    1.30x baseline
paddedFma     1.237 ns/op    1.31x baseline

but it’s interesting I remember that the paddedFma used to win by more than 15%