c - C中矩阵和 vector 乘法的优化

Question

我有一个函数，它获取一个 3 x 3 矩阵和一个 3 x 4000 vector，并将它们相乘。 所有计算均以 double (64 位)完成。 该函数被调用了大约 3.5 百万次，因此应该对其进行优化。

#define MATRIX_DIM 3
#define VECTOR_LEN 3000

typedef struct {
    double a;
    double b;
    double c;
} vector_st;

double matrix[MATRIX_DIM][MATRIX_DIM];
vector_st vector[VACTOR_LEN];

inline void rotate_arr(double input_matrix[][MATRIX_DIM], vector_st *input_vector, vector_st *output_vector)
{
    int i;
    for (i = 0; i < VACTOR_LEN; i++) {
        op_rotate_preset_arr[i].a = input_matrix[0][0] * input_vector[i].a + 
                                    input_matrix[0][1] * input_vector[i].b +
                                    input_matrix[0][2] * input_vector[i].c;

        op_rotate_preset_arr[i].b = input_matrix[1][0] * input_vector[i].a + 
                                    input_matrix[1][1] * input_vector[i].b +
                                    input_matrix[1][2] * input_vector[i].c;

        op_rotate_preset_arr[i].c = input_matrix[2][0] * input_vector[i].a + 
                                    input_matrix[2][1] * input_vector[i].b +
                                    input_matrix[2][2] * input_vector[i].c;
    }
}

我完全没有关于如何优化它的想法，因为它是内联，数据访问是顺序的，而且功能很短而且非常简单。 可以假设 vector 始终相同，只有矩阵在变化才能提高性能。

Answer 1

这里一个容易解决的问题是编译器假设矩阵和输出 vector 可能混叠。如所见here在第二个函数中，这会导致生成的代码效率较低但代码量大得多。这可以简单地通过将 restrict 添加到输出指针来解决。仅这样做已经有助于并使代码不受平台特定优化的影响，但依赖于自动矢量化以利用过去二十年发生的性能提升。

自动矢量化对于这项任务显然还不够成熟，Clang 和 GCC 都会产生太多的数据混洗。这应该会在未来的编译器中得到改进，但现在即使是这样的情况(看起来并不是天生的 super 难)也需要手动帮助，比如这个(虽然没有测试)

void rotate_arr_avx(double input_matrix[][MATRIX_DIM], vector_st *input_vector, vector_st * restrict output_vector)
{
    __m256d col0, col1, col2, a, b, c, t;
    int i;
    // using set macros like this is kind of dirty, but it's outside the loop anyway
    col0 = _mm256_set_pd(0.0, input_matrix[2][0], input_matrix[1][0], input_matrix[0][0]);
    col1 = _mm256_set_pd(0.0, input_matrix[2][1], input_matrix[1][1], input_matrix[0][1]);
    col2 = _mm256_set_pd(0.0, input_matrix[2][2], input_matrix[1][2], input_matrix[0][2]);
    for (i = 0; i < VECTOR_LEN; i++) {
        a = _mm256_set1_pd(input_vector[i].a);
        b = _mm256_set1_pd(input_vector[i].b);
        c = _mm256_set1_pd(input_vector[i].c);
        t = _mm256_add_pd(_mm256_add_pd(_mm256_mul_pd(col0, a), _mm256_mul_pd(col1, b)), _mm256_mul_pd(col2, c));
        // this stores an element too much, ensure 8 bytes of padding exist after the array
        _mm256_storeu_pd(&output_vector[i].a, t);
    }
}

以这种方式编写它可以显着改进编译器对它的处理，现在编译成一个漂亮而紧凑的循环，没有所有废话。早些时候的代码看起来很难看，但是有了这个循环现在看起来像这样(GCC 8.1，启用了 FMA)，它实际上是可读的:

.L2:
    vbroadcastsd    ymm2, QWORD PTR [rsi+8+rax]
    vbroadcastsd    ymm1, QWORD PTR [rsi+16+rax]
    vbroadcastsd    ymm0, QWORD PTR [rsi+rax]
    vmulpd  ymm2, ymm2, ymm4
    vfmadd132pd     ymm1, ymm2, ymm3
    vfmadd132pd     ymm0, ymm1, ymm5
    vmovupd YMMWORD PTR [rdx+rax], ymm0
    add     rax, 24
    cmp     rax, 72000
    jne     .L2

这有一个明显的不足:实际使用了 256 位 AVX vector 的 4 个 double 槽中的 3 个。如果 vector 的数据格式更改为例如 AAAABBBBCCCC 重复，则可以使用完全不同的方法，即广播矩阵元素而不是 vector 元素，然后将广播的矩阵元素乘以 4 个不同 的 A 分量vector_st一次。

我们可以尝试的另一件事是在不改变数据格式的情况下同时处理多个矩阵，这有助于重新使用来自 input_vector 的负载以增加算术强度。

void rotate_arr_avx(double input_matrixA[][MATRIX_DIM], double input_matrixB[][MATRIX_DIM], vector_st *input_vector, vector_st * restrict output_vectorA, vector_st * restrict output_vectorB)
{
    __m256d col0A, col1A, col2A, a, b, c, t, col0B, col1B, col2B;
    int i;
    // using set macros like this is kind of dirty, but it's outside the loop anyway
    col0A = _mm256_set_pd(0.0, input_matrixA[2][0], input_matrixA[1][0], input_matrixA[0][0]);
    col1A = _mm256_set_pd(0.0, input_matrixA[2][1], input_matrixA[1][1], input_matrixA[0][1]);
    col2A = _mm256_set_pd(0.0, input_matrixA[2][2], input_matrixA[1][2], input_matrixA[0][2]);
    col0B = _mm256_set_pd(0.0, input_matrixB[2][0], input_matrixB[1][0], input_matrixB[0][0]);
    col1B = _mm256_set_pd(0.0, input_matrixB[2][1], input_matrixB[1][1], input_matrixB[0][1]);
    col2B = _mm256_set_pd(0.0, input_matrixB[2][2], input_matrixB[1][2], input_matrixB[0][2]);
    for (i = 0; i < VECTOR_LEN; i++) {
        a = _mm256_set1_pd(input_vector[i].a);
        b = _mm256_set1_pd(input_vector[i].b);
        c = _mm256_set1_pd(input_vector[i].c);
        t = _mm256_add_pd(_mm256_add_pd(_mm256_mul_pd(col0A, a), _mm256_mul_pd(col1A, b)), _mm256_mul_pd(col2A, c));
        // this stores an element too much, ensure 8 bytes of padding exist after the array
        _mm256_storeu_pd(&output_vectorA[i].a, t);
        t = _mm256_add_pd(_mm256_add_pd(_mm256_mul_pd(col0B, a), _mm256_mul_pd(col1B, b)), _mm256_mul_pd(col2B, c));
        _mm256_storeu_pd(&output_vectorB[i].a, t);
    }
}