SIMD product of a 4×4 matrix and a vector

Question

I wrote a matrix class for use in a ray tracer. I already have a 'scalar' version up and running, but now I'm trying to rewrite it to use Intel SIMD Instrinsics. I realize my compiler (clang-7.0.0) will probably generate code that is at least as fast as mine (chances are, it's going to be faster), but I'm doing this because I enjoy it, and because I find it interesting to think about things like this.

This is my version of a matrix-vector product. The matrix has a fixed dimension of 4x4, the vector has a fixed dimension of 4. I was wondering if I could take a different approach to make this function even faster?

My matrix class is set-up like this:

class alignas(16 * sizeof(float)) matrix
{
    private:
        union
        {
            struct
            {
                __m256 ymm0;
                __m256 ymm1;
            };

            float data [16];
        };

        ...

    public:
        ...
}

ymm0 stores rows 0 and 1, ymm1 holds rows 2 and 3.

My simd_vec4 class is just contains a single __m128 called xmm that's aligned to 16 bytes (4 * sizeof(float)). The vector stores its components in the as {x, y, z, w}.

Here's my reasoning for the algorithm:

1. We basically need to perform 4 dot products; so the first thing I do is broadcast my 4-float vector to an 8-float vector with the same 4 floats in every lane.

2. Now, we perform the necessary multiplications, these only carry a data dependency with the vector from step 1, so they can be parallelized/out-of-order'ed.

3. (a) The next step is to reduce our 16 floats down to 4 floats. We do this by first performing 1 horzontal reduction on each of our 2-lane vectors, and then swapping lanes in one of them.

   (b) The second portion of our reduction operation is to merge our two 2-lane vectors into one 2-lane vector, due to the way our data is laid out, we get the same data in both lanes (different order, but we need to fix that anyway...).

4. We now extract only the lower lane, because it has all the data we need.

5. Finally, we shuffle some data around in the lower lane to put the right components in the right places.

And here's the code:

simd_vec4 operator*(const simd_mat4& lhs, const simd_vec4& rhs) noexcept
{
                                                                        // ymm0    = [        A                 B                 C                 D        ][        E                 F                 G                 H        ]
                                                                        // ymm1    = [        I                 J                 K                 L        ][        M                 N                 O                 P        ]
    __m256 broadcast = _mm256_set_m128(rhs.xmm, rhs.xmm);               // Vec     = [         x                 y                 z                 w       ][         x                 y                 z                 w       ]
                                                                        //-------------------------------------------------------------------------------------------------------------------------------------------------------------
    __m256 xy_m256 = _mm256_mul_ps(lhs.ymm0, broadcast);                // xy_m256 = [        Ax                By                Cz                Dw       ][        Ex                Fy                Gz                Hw       ]
    __m256 zw_m256 = _mm256_mul_ps(lhs.ymm1, broadcast);                // zw_m256 = [        Ix                Jy                Kz                Lw       ][        Mx                Ny                Oz                Pw       ]
                                                                        //-------------------------------------------------------------------------------------------------------------------------------------------------------------
    __m256   acc0  = _mm256_hadd_ps(xy_m256, zw_m256);                  // acc0    = [     Ix + Jy           Kz + Lw           Ax + By           Cz + Dw     ][     Mx + Ny           Oz + Pw           Ex + Fy           Gz + Hw     ]
    __m256   acc1  = _mm256_hadd_ps(zw_m256, xy_m256);                  // acc1    = [     Ax + By           Cz + Dw           Ix + Jy           Kz + Lw     ][     Ex + Fy           Gz + Hw           Mx + Ny           Oz + Pw     ]
             acc1  = _mm256_permute2f128_ps(acc1, acc1, 0x01);          // acc1    = [     Ex + Fy           Gz + Hw           Mx + Ny           Oz + Pw     ][     Ax + By           Cz + Dw           Ix + Jy           Kz + Lw     ]
    __m256 merged  = _mm256_hadd_ps(acc0, acc1);                        // merged  = [Ex + Fy + Gz + Hw Mx + Ny + Oz + Pw Ix + Jy + Kz + Lw Ax + By + Cz + Dw][Ax + By + Cz + Dw Ix + Jy + Kz + Lw Mx + Ny + Oz + Pw Ex + Fy + Gz + Hw]
    __m128    vec  = _mm256_extractf128_ps(merged, 0);                  //  vec    =                                                                          [Ax + By + Cz + Dw Ix + Jy + Kz + Lw Mx + Ny + Oz + Pw Ex + Fy + Gz + Hw]
              vec  = _mm_shuffle_ps(vec, vec, _MM_SHUFFLE(2, 1, 3, 0)); //  vec    =                                                                          [Mx + Ny + Oz + Pw Ix + Jy + Kz + Lw Ex + Fy + Gz + Hw Ax + By + Cz + Dw]

    return simd_vec4(vec);
}

I'm running this code on an Intel Core i7-5500U CPU @ 2.40GHz, so the allowed extensions are SSE4.1, SSE4.2, AVX2 (and below, I'm assuming).

Any advice is welcome, but one thing I do struggle with is coming up with meaningful names for the intermediate results, if anyone has any tips, you're more than welcome!

I've written two unit test so far, and both are passing (both are 'normal' cases though, I'm working on adding more 'edge' cases)

Here's the unit test code (I use Catch2 btw):

TEST_CASE("matrix * vec4", "[simd_matrix][vec4][multiplication][product]")
{
    Math::matrix mat;
    Math::simd_vec4 vec;

    // First test case, tested by writing out all operations
    constexpr std::array<float, 16> TEST_CASE_1 =
                                        {0,  1,  2,  3,
                                         4,  5,  6,  7,
                                         8,  9,  10, 11,
                                         12, 13, 14, 15};

    for (int i = 0; i < 4; i++)
    {
        for (int j = 0; j < 4; j++)
        {
            const int arrayIdx = (i * 4) + j;

            mat(i, j) = TEST_CASE_1[arrayIdx];
        }
    }

    vec = Math::simd_vec4(1.0f, 3.0f, 2.0f, 5.0f);

    Math::simd_vec4 result   = mat * vec;
    Math::simd_vec4 EXPECTED_1 = Math::simd_vec4(
        (0.0f * 1.0f)  + (1.0f * 3.0f)  + (2.0f * 2.0f)  + (3.0f * 5.0f),
        (4.0f * 1.0f)  + (5.0f * 3.0f)  + (6.0f * 2.0f)  + (7.0f * 5.0f),
        (8.0f * 1.0f)  + (9.0f * 3.0f)  + (10.0f * 2.0f) + (11.0f * 5.0f),
        (12.0f * 1.0f) + (13.0f * 3.0f) + (14.0f * 2.0f) + (15.0f * 5.0f));

    CHECK(result.x == Approx(EXPECTED_1.x));
    CHECK(result.y == Approx(EXPECTED_1.y));
    CHECK(result.z == Approx(EXPECTED_1.z));
    CHECK(result.w == Approx(EXPECTED_1.w));

    // Second test case, checked with Wolfram Alpha (https://www.wolframalpha.com/input/?i=%7B%7B1,+7,+23,+-5%7D,%7B0,+-9,+-5,+1%7D,%7B2,+6,+-3,+8%7D,%7B-1,+8,+11,+-5%7D%7D+*+%7B-7,+5,+-3,+1%7D)
    constexpr std::array<float, 16> TEST_CASE_2 =
                                        {1, 7, 23, -5,
                                         0, -9, -5, 1,
                                         2, 6, -3, 8,
                                         -1, 8, 11, -5};

    for (int i = 0; i < 4; i++)
    {
        for (int j = 0; j < 4; j++)
        {
            const int arrayIdx = (i * 4) + j;
            mat(i, j) = TEST_CASE_2[arrayIdx];
        }
    }

    vec = Math::simd_vec4(-7, 5, -3, 1);
    Math::simd_vec4 EXPECTED_2 = Math::simd_vec4(-46, -29, 33, 9);
    result   = mat * vec;

    CHECK(result.x == Approx(EXPECTED_2.x));
    CHECK(result.y == Approx(EXPECTED_2.y));
    CHECK(result.z == Approx(EXPECTED_2.z));
    CHECK(result.w == Approx(EXPECTED_2.w));
}

For those that are interested: The full code can be found here: Header, Implementation. Should you have any problems building, here's my CMakeLists.txt.

Edit: I quickly got a gist setup, it's got 5 files (2 for the matrix, 2 for the vector and 1 for main) along with a shell script that contains the command I used to compile this.

The example code runs the first test case. If any of the results are wrong, the error will be written to stderr.

Answer 1

An issue with that code is that while it tries to pack the data to make efficient use of arithmetic throughput, actually arithmetic throughput is high anyway and it's the shuffles (including horizontal addition which has two shuffles internally) that are relatively expensive. Shuffles don't all have a low latency either, cross-slice shuffles such as vperm2f128 take 3 cycles, so it's easy to accidentally build up a large delay that way.

So, as far as I know, efficient single-vector mat4 x vec4 is still based on broadcasting the elements of the vector, multiplying that by columns of the matrix, and adding up the results (compilers tend to merge the add/mul into FMA if allowed). This would cost 4 shuffles pre-AVX or if the vector comes in a register, and potentially zero shuffles with AVX if the vector comes from memory, since a broadcast-from-memory has a "free shuffle" (it's just a load, and broadcasts for free, rather than going to the shuffle unit, though pre-AVX512 a broadcast-load cannot be a memory operand to an arithmetic operation). For example:

__m128 transform4(__m128* mat4, float *vec4) {
    __m128 x = _mm_set1_ps(vec4[0]);
    __m128 y = _mm_set1_ps(vec4[1]);
    __m128 z = _mm_set1_ps(vec4[2]);
    __m128 w = _mm_set1_ps(vec4[3]);

    __m128 p1 = _mm_mul_ps(x, mat4[0]);
    __m128 p2 = _mm_mul_ps(y, mat4[1]);
    __m128 p3 = _mm_mul_ps(z, mat4[2]);
    __m128 p4 = _mm_mul_ps(w, mat4[3]);

    return _mm_add_ps(_mm_add_ps(p1, p2), _mm_add_ps(p3, p4));
}

With Clang 7, AVX2 enabled:

transform4(float __vector(4)*, float*):                 # @transform4(float __vector(4)*, float*)
    vbroadcastss    xmm0, dword ptr [rsi]
    vbroadcastss    xmm1, dword ptr [rsi + 4]
    vbroadcastss    xmm2, dword ptr [rsi + 8]
    vbroadcastss    xmm3, dword ptr [rsi + 12]
    vmulps  xmm0, xmm0, xmmword ptr [rdi]
    vfmadd231ps     xmm0, xmm1, xmmword ptr [rdi + 16] # xmm0 = (xmm1 * mem) + xmm0
    vfmadd231ps     xmm0, xmm2, xmmword ptr [rdi + 32] # xmm0 = (xmm2 * mem) + xmm0
    vfmadd231ps     xmm0, xmm3, xmmword ptr [rdi + 48] # xmm0 = (xmm3 * mem) + xmm0
    ret

This has a decent throughput, one transform every four cycles in the best case as it is, and better in a loop with the loads from the matrix factored out of the loop (hopefully a compiler can do that but better check the asm to make sure). The FMAs are tied together though (by Clang!), so the latency is not the greatest. Therefore it is best to batch transforms as much as reasonably possible, or otherwise latency could be reduced at the cost of an extra addition (and loss of source-level compatibility for pre-FMA processors)

__m128 transform4(__m128* mat4, float *vec4) {
    __m128 x = _mm_set1_ps(vec4[0]);
    __m128 y = _mm_set1_ps(vec4[1]);
    __m128 z = _mm_set1_ps(vec4[2]);
    __m128 w = _mm_set1_ps(vec4[3]);

    __m128 p1 = _mm_mul_ps(x, mat4[0]);
    __m128 p2 = _mm_fmadd_ps(y, mat4[1], p1);
    __m128 p3 = _mm_mul_ps(z, mat4[2]);
    __m128 p4 = _mm_fmadd_ps(w, mat4[3], p3);

    return _mm_add_ps(p2, p4);
}

These approaches don't scale well to wider SIMD..

With 8 SoA vectors to work with though, we can switch to broadcasting the matrix elements and doing 16 multiplications, which halves the number arithmetic instructions per transform. Since the matrix is being broadcasted from, its format is now free to choose, while this time the vector format is constrained. For example (not tested):

void transformBatch8(float *mat4, float *v) {
    __m256 m00 = _mm256_set1_ps(mat4[0]);
    __m256 m01 = _mm256_set1_ps(mat4[1]);
    __m256 m02 = _mm256_set1_ps(mat4[2]);
    __m256 m03 = _mm256_set1_ps(mat4[3]);
    __m256 m10 = _mm256_set1_ps(mat4[4]);
    __m256 m11 = _mm256_set1_ps(mat4[5]);
    __m256 m12 = _mm256_set1_ps(mat4[6]);
    __m256 m13 = _mm256_set1_ps(mat4[7]);
    __m256 m20 = _mm256_set1_ps(mat4[8]);
    __m256 m21 = _mm256_set1_ps(mat4[9]);
    __m256 m22 = _mm256_set1_ps(mat4[10]);
    __m256 m23 = _mm256_set1_ps(mat4[11]);
    __m256 m30 = _mm256_set1_ps(mat4[12]);
    __m256 m31 = _mm256_set1_ps(mat4[13]);
    __m256 m32 = _mm256_set1_ps(mat4[14]);
    __m256 m33 = _mm256_set1_ps(mat4[15]);
    __m256 x = _mm256_load_ps(v);
    __m256 y = _mm256_load_ps(v + 8);
    __m256 z = _mm256_load_ps(v + 16);
    __m256 w = _mm256_load_ps(v + 24);
    __m256 rx = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_mul_ps(m00, x), m01, y), m02, z), m03, w);
    __m256 ry = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_mul_ps(m10, x), m11, y), m12, z), m13, w);
    __m256 rz = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_mul_ps(m20, x), m21, y), m22, z), m23, w);
    __m256 rw = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_mul_ps(m30, x), m31, y), m32, z), m33, w);
    _mm256_store_ps(v, rx);
    _mm256_store_ps(v + 8, ry);
    _mm256_store_ps(v + 16, rz);
    _mm256_store_ps(v + 24, rw);
}

Based on the cost of 20 loads, it takes at least 10 cycles per 8 transforms this way. Putting this in a loop and factoring out some loads would help, not all loads need to be factored out (and pre-AVX512 they cannot be, there are not enough registers), at least 4 so that the total number of loads per batch of 8 transforms goes down to 16, bringing down the minimum time down to 8 cycles per 8 transforms (in practice it often helps to reduce loads even further, but only to get closer to that bound of 8 cycles per 8 transforms because the FMAs alone already enforce that lower limit).