OpenMP实现三角矩阵向量乘法

271 阅读2分钟

开启掘金成长之旅!这是我参与「掘金日新计划 · 12 月更文挑战」的第21天,点击查看活动详情

矩阵向量乘

矩阵乘操作是很多应用领域中的关键内核,比如高性能计算,人工智能和大数据。在人工智能中分类模型在最后都要从特征矩阵得到一个向量(一般是在线性层完成),这个便是矩阵向量乘操作。

给定矩阵例如

A = [[1, 1, 1, 1],  
     [0, 1, 1, 1],
     [0, 0, 1, 1],         b = [1, 1, 1, 1], 则A × b = [4, 3, 2, 1]
     [0, 0, 0, 1]
    ]

本文要求是下三角,并且矩阵按列存储。

image.png

注意一下流程:

  1. malloc数据
  2. 矩阵和向量初始化(简单要求下矩阵是列主序的)
  3. 写原始的for循环完成矩阵向量乘
  4. 加上omp制导语句,改写for循环
#include <stdio.h>
#include <omp.h>
#include <stdlib.h>

int main(int argc, char** argv) {
    int m = 6;     //矩阵大小 m*m

    int nthreads =  6;
    omp_set_num_threads(nthreads);
    int tid = omp_get_num_threads();

    int* A = (int*)malloc(sizeof(int)*m*m);
    int* b = (int*)malloc(sizeof(int)*m);
    int* c = (int*)malloc(sizeof(int)*m);
    int i, j;
    //上三角矩阵初始化, 列主序
    #pragma omp parallel for collapse(2) private(j) shared(A, m)
    for(i=0; i<m; i++) {
        for(j=0; j<m; j++) {
            if(i>=j) {
                A[i*m+j] = 1;
            } else {
                A[i*m+j] = 0;
            }
        }
    }

    int *pos = A;
    for (int i = 0; i < m; i++){
        for(int j=0; j<m; j++)
            printf("%d ", *(pos++));
        printf("\n");
    }
    

    //向量初始化
    for(int i=0; i<m; i++) {
        b[i] = 1;
        c[i] = 0;
    }

    //创建一个互斥锁
    omp_lock_t lock;
    omp_init_lock(&lock);

    //矩阵向量乘法
    #pragma omp parallel for collapse(2) private(j) shared(A, b, c, m)
    for(int i=0; i<m; i++) {
        for(int j=0; j<m; j++) {
            omp_set_lock(&lock);
            c[j] += A[i*m+j] * b[j];    //c[j]的写操作是互斥的,否则输出结果不正确
            omp_unset_lock(&lock);
        }
    } 

    //输出结果
    for(int i=0; i<m; i++) {
        printf("%d ", c[i]);
    }
}

由于是按照列存储,所以c[j] += A[i*m+j] * b[j],如果是按行存储,则c[i] += A[i*m+j] * b[j]。 由于是按照列存储,在按照指针遍历时,会发现上三角变成了下三角。

运行结果:

1 0 0 0 0 0   //这是按列存储之后,以指针方式遍历矩阵
1 1 0 0 0 0
1 1 1 0 0 0
1 1 1 1 0 0
1 1 1 1 1 0
1 1 1 1 1 1

6 5 4 3 2 1   //这是相乘之后的结果