如何通过寄存器数据复用加速矩阵向量间Pearson相关系数计算

背景 #

我正在学习并行编程课程，核心知识点是处理器寄存器内数据复用技术（扩展至缓存层级）——即在从内存加载新数据前完成尽可能多的计算，课程使用C++语言。

我尝试将该技术应用于以下问题：给定包含向量$x_i$的矩阵$M$，计算矩阵中每对向量间的Pearson相关系数，结果存储至矩阵$R$，因系数对称仅需计算上三角部分。

朴素实现版本 #

最初采用双层循环遍历矩阵行，先对矩阵预处理，利用向量类型并行执行乘法与求和。计算系数需获取向量点积及元素和，代码如下：

#include <cmath>
#include <vector>

using namespace std;

#define PARTITION_SIZE 8
#define INIT_PARTITION {.0,.0,.0,.0,.0,.0,.0,.0}

typedef double double4_t __attribute__ ((vector_size (PARTITION_SIZE * sizeof(double))));

void correlate(int ny, int nx, const float *data, float *result) {

    int partitions_per_row = nx % PARTITION_SIZE != 0 ? nx/PARTITION_SIZE + 1 : nx/PARTITION_SIZE;

    vector<double4_t> processed_data((partitions_per_row) * ny);

    /* Load data into vectorized format for easier processing */
    // Traverse rows
    for(int i = 0; i < ny; i++){
        // Traverse partitions in row
        for(int j = 0; j < partitions_per_row; j++){
            // Traverse elements inside a partition
            for (int k = 0; k < PARTITION_SIZE; ++k) {
                int old_x_index = j * PARTITION_SIZE + k;
                // If the element is inside the matrix, load it, otherwise load 0 for padding
                // This way we can have an integer number of partitions per row
                processed_data[i * partitions_per_row + j][k] = old_x_index < nx ? (data[i * nx + old_x_index] ): 0.0;
            }
        }
    }

    /* Calculate Pearson's coefficient between each vector of the matrix */
    for(int i = 0; i < ny; i++){
        double4_t x_partition;
        double4_t y_partition;  

        for(int j = i; j < ny; j++){
            double4_t sum_of_x_partition = INIT_PARTITION;
            double4_t sum_of_y_partition = INIT_PARTITION;
            double4_t sum_of_x_squared_partition = INIT_PARTITION;
            double4_t sum_of_y_squared_partition = INIT_PARTITION;
            double4_t sum_of_x_times_y_partition = INIT_PARTITION;

            int x_row_index = i * partitions_per_row;
            int y_row_index = j * partitions_per_row;
            
            // Traverse partitions in row
            for(int k = 0; k < partitions_per_row; k++){
                x_partition = processed_data[x_row_index + k];
                y_partition = processed_data[y_row_index + k];

                // Calculate the sums of each partition
                sum_of_x_partition += x_partition;
                sum_of_y_partition += y_partition;
                sum_of_x_squared_partition += x_partition * x_partition;
                sum_of_y_squared_partition += y_partition * y_partition;
                sum_of_x_times_y_partition += x_partition * y_partition;
            }

            // Agregate the different sums
            for (int k = 1; k < PARTITION_SIZE; ++k) {
                sum_of_x_partition[0] += sum_of_x_partition[k];
                sum_of_y_partition[0] += sum_of_y_partition[k];
                sum_of_x_squared_partition[0] += sum_of_x_squared_partition[k];
                sum_of_y_squared_partition[0] += sum_of_y_squared_partition[k];
                sum_of_x_times_y_partition[0] += sum_of_x_times_y_partition[k];
            }

            // Calculate the correlation
            double numerator = nx * sum_of_x_times_y_partition[0] - sum_of_x_partition[0] * sum_of_y_partition[0];
            double denominator = sqrt((nx * sum_of_x_squared_partition[0] - sum_of_x_partition[0] * sum_of_x_partition[0]) * 
                                                        (nx * sum_of_y_squared_partition[0] - sum_of_y_partition[0] * sum_of_y_partition[0]));

            double correlation = numerator / denominator;

            result[i * ny + j] = correlation;
        }
    }
}

寄存器复用的首次实现 #

我发现常规遍历难以复用数据，故先转置矩阵，使每行包含各向量的对应元素，再转换为向量格式以利用处理器向量操作。思路是通过整向量乘法结合元素置换获取所有元素组合，减少运算量，同时重复遍历同一行提升读写效率。使用_mm256_permute4x64_pd与_mm256_permute_pd置换向量，乘积存入temporal_mults，向量和存入sums_vector。

但该实现虽可运行，性能却远低于朴素版本，无法确定是实现缺陷还是技术理解有误，特附上代码：

#include <cmath>
#include <vector>
#include <x86intrin.h>

using namespace std;

// The registers are 256 bits long, so we can fit 4 doubles in one register.
#define PARTITION_SIZE 4
#define INIT_PARTITION {.0,.0,.0,.0}

typedef double double4_t __attribute__ ((vector_size (PARTITION_SIZE * sizeof(double))));

static inline double4_t swap2(double4_t x) { return _mm256_permute4x64_pd(x, 0b01001110); }
static inline double4_t swap1(double4_t x) { return _mm256_permute_pd(x, 0b0101); }


/* Load data into vectorized format for easier processing.
   It is a bit hard to wrap your head around the indeces,
   but basically we are creating the transpose of the matrix, 
   vectorizing the data and adding padding at the same time.
*/
vector<double4_t> optimizeMatrix(int nx, int ny, int partitions_per_column, const float *data){
    vector<double4_t> processed_data(nx*partitions_per_column);
    // Traverse columns
    for(int i = 0; i < nx; i++){
        // Traverse rows
        for(int j = 0; j < partitions_per_column; j++){
            int processed_data_column = i * partitions_per_column;
            int processed_vector_index = processed_data_column + j;

            for (int k = 0; k < PARTITION_SIZE; ++k) {
                int old_column = j * PARTITION_SIZE + k;
                processed_data[processed_vector_index][k] = 
                    old_column < ny ? data[old_column * nx + i] : .0; 
            }
        }
    }
    return processed_data;
}

void computeRowMultiplications(
        int ny,
        int processed_data_row,
        int partitions_per_column,
        vector<double4_t> &processed_data,
        vector<double4_t> &sums_vector,
        vector<double> &temporal_mults){

    for(int partition_a = 0; partition_a < partitions_per_column; partition_a++){
        double4_t a000 = processed_data[processed_data_row + partition_a];
        double4_t a001 = swap1(a000);

        sums_vector[partition_a] += a000;

        int starting_ele_index = partition_a * PARTITION_SIZE; // Translate partition number to vector index
        int result_row_1 = starting_ele_index * ny; // Calculate result row for first element from vector index
        
        for(int partition_b = partition_a; partition_b < partitions_per_column; partition_b++){
            double4_t b000 = processed_data[processed_data_row + partition_b];
            double4_t b001 = swap1(b000);
            double4_t b010 = swap2(b000);

            double4_t mult_vector[PARTITION_SIZE];
            mult_vector[0] = a000 * b000;
            mult_vector[1] = a000 * b001;
            mult_vector[2] = a000 * b010;
            mult_vector[3] = swap1(a001 * b010);

            short contador_2 = -1, contador_4 = -2;
            for(int perm = 0, perm_opposite = PARTITION_SIZE - 1; perm < PARTITION_SIZE; perm++, perm_opposite--){
                contador_2 *= -1;
                if (perm % 2 == 0) {
                    contador_4 *= -1;
                }
                temporal_mults[result_row_1 + partition_b * PARTITION_SIZE + perm] += mult_vector[perm][0];
                temporal_mults[result_row_1 + ny + partition_b * PARTITION_SIZE + perm] += mult_vector[perm + contador_2][1];
                temporal_mults[result_row_1 + ny + ny + partition_b * PARTITION_SIZE + perm] += mult_vector[perm + contador_4][2];
                temporal_mults[result_row_1 + ny + ny + ny + partition_b * PARTITION_SIZE + perm] += mult_vector[perm_opposite][3];
            }
        }
    }
}

void correlate(int ny, int nx, const float *data, float *result) {

    int partitions_per_column = ny % PARTITION_SIZE != 0 ? ny/PARTITION_SIZE + 1 : ny/PARTITION_SIZE;

    vector<double4_t> processed_data;
    vector<double> temporal_mults( pow(partitions_per_column * PARTITION_SIZE, 2), 0.0);

    vector<double4_t> sums_vector(partitions_per_column);
    for (auto& vec : sums_vector) {
        vec[0] = 0.0;
        vec[1] = 0.0;
        vec[2] = 0.0;
        vec[3] = 0.0;
    }
 
    processed_data = optimizeMatrix(nx, ny, partitions_per_column, data);

    /* Calculate the multiplication between vectors and store them in temporal_muls,
       and the element sum of each vector and store them in sums_vectors */
    for(int i = 0; i < nx; i++){
        int current_row = i*partitions_per_column;
        computeRowMultiplications(
            ny,
            current_row, 
            partitions_per_column, 
            processed_data,
            sums_vector,
            temporal_mults);

    }

    /* Use the values computed above to calculate the correlations between vectors,
       and store them in results. */
    for(int i = 0; i < ny; i++){
        int current_row = i * ny;

        double sum_i = sums_vector[i/PARTITION_SIZE][i % PARTITION_SIZE];
        double squared_sum_i = temporal_mults[current_row + i];
        for(int j = ny-1; j >= i; j--){
            double sum_j = sums_vector[j/PARTITION_SIZE][j % PARTITION_SIZE];
            double squared_sum_j = temporal_mults[j*ny + j];

            int current_cell = current_row + j;
            double sum_of_i_times_j = temporal_mults[current_cell];

            double numerator = nx * sum_of_i_times_j - sum_i * sum_j;
            double denominator = sqrt((nx * squared_sum_i - sum_i * sum_i) * 
                                      (nx * squared_sum_j - sum_j * sum_j));

            result[current_cell] = numerator / denominator;
        }
    }
}

求助 #

作为并行高效编程新手，我确信存在诸多未察觉的问题，恳请各位提供帮助。若问题表述存在不当之处，也请指出。

内容的提问来源于stack exchange，提问作者Lucas Piorno