4×4矩阵和向量的SIMD乘积

Question

我编写了一个矩阵类，用于射线追踪器。我已经启动并运行了一个“标量”版本，但现在我正在尝试重写它以使用Intel SIMD Instrinsics。我意识到我的编译器(clang-7.0.0)可能会生成至少和我的编译器一样快的代码(很有可能，它会更快)，但是我这样做是因为我喜欢它，因为我发现思考这样的事情很有趣。

这是我的矩阵向量积的版本。矩阵的固定维数为4x4，向量的固定维数为4，我想知道是否可以采用不同的方法使这个函数更快？

我的矩阵类是这样设置的：

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

            float data [16];
        };

        ...

    public:
        ...
}

ymm0存储行0和1，ymm1存储行2和3。

我的simd_vec4类只包含一个名为xmm的__m128，它与16个字节(4 * sizeof(float))对齐。向量将其组件存储在as {x, y, z, w}中。

下面是我对算法的推理：

1. 我们基本上需要执行4点产品，所以我做的第一件事是将我的4-浮点数矢量广播到8-浮动向量，在每个车道上都有相同的4种浮动。
2. 现在，我们执行必要的乘法，这些乘法只带有步骤1中向量的数据依赖项，因此它们可以并行/无序。
3. (a)下一步是将我们的16辆车减少到4辆。我们首先在我们的双线向量上执行一个水平减少，然后在其中一个中交换车道。(b)我们减缩操作的第二部分是将两个双线向量合并成一个双线向量，由于数据的布局方式，我们在两条车道上得到相同的数据(不同的顺序，但无论如何我们需要修正.)。
4. 我们现在只提取下车道，因为它有我们需要的所有数据。
5. 最后，我们在较低的车道上对一些数据进行洗牌，以便将正确的部件放在正确的位置。

这是密码：

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);
}

我是在IntelCorei7-5500Ucpu@ 2.40GHz上运行这段代码的，所以允许的扩展是SSE4.1、SSE4.2、AVX2 2 (我假设在下面)。

任何建议都是受欢迎的，但有一件事我很难接受，那就是为中间结果找出有意义的名字，如果有人有任何建议的话，你就更受欢迎了！

到目前为止，我已经编写了两个单元测试，它们都通过了(但都是“正常”案例，我正在努力添加更多的“边缘”案例)

下面是单元测试代码(我使用Catch2 )：

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));
}

对于那些感兴趣的人，可以在这里找到完整的代码：标题，实现。如果你有任何问题，建设，这是我的CMakeLists.txt。

编辑:我很快就得到了一个要旨设置，它有5个文件(2个用于矩阵，2个用于向量，1个用于main)，以及一个包含用于编译此命令的shell脚本。

示例代码运行第一个测试用例。如果任何结果都是错误的，则错误将写入stderr。

Answer 1

该代码的一个问题是，尽管它试图打包数据以有效地利用算术吞吐量，但实际上算术吞吐量还是很高的，而且相对昂贵的是洗牌(包括内部有两个洗牌的水平加法)。洗牌并不是所有的延迟都很低，像vperm2f128这样的跨片洗牌需要3个周期，所以很容易意外地造成很大的延迟。

因此，据我所知，高效的单矢量mat4 x vec4仍然基于广播向量的元素，将其乘以矩阵的列，并将结果相加(编译器倾向于将add/mul合并到FMA中，如果允许的话)。这将花费4次预AVX或向量进入寄存器中，如果矢量来自内存，则可能为零，因为从内存广播有“免费洗牌”(这只是一种负载，广播是免费的，而不是进入洗牌单元，尽管AVX 512前的广播不能作为内存操作数进行算术操作)。例如：

__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));
}

在Clang 7下，启用了AVX2：

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

这具有很好的吞吐量，在最好的情况下每四个周期进行一次转换，在循环中更好地使用从循环中提取的矩阵负载(希望编译器能够做到这一点，但更好地检查asm以确保)。虽然FMA是绑定在一起的(由Clang!)，所以延迟并不是最大的。因此，最好是尽可能多地批处理转换，否则就可以减少延迟，而代价是额外的添加(以及失去前FMA处理器的源级兼容性)。

__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);
}

这些方法不能很好地扩展到更广泛的SIMD。

使用8个SoA向量，我们可以切换到广播矩阵元素，进行16次乘法，使每个变换的数字算术指令减半。由于矩阵是广播的，它的格式现在可以自由选择，而这一次矢量格式受到限制。例如(未测试)：

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);
}

基于20种负载的成本，这种方式每8次转换至少需要10个周期。将其放入循环并分解一些负载将有帮助，并不是所有的负载都需要被分解(而且预and 512不能，没有足够的寄存器)，至少有4个，这样每批8次转换的总负载将降至16，使最小时间降至每8次转换8周期(在实践中，它通常有助于进一步减少负载，但只是接近于8/8转换的8周期的界限，因为FMA本身已经执行了这一下限)。