如何使用 OpenMP 并行化 COO 格式的稀疏矩阵-矩阵 (SpMM) 乘积？

Question

我最近使用 COO 存储格式编写了稀疏矩阵向量 (SpMV) 产品的多线程版本：

void dcoomv(SparseMatrixCOO *A, double *x, double *y) {
  for (int i=0; i<A->m; i++) y[i] = 0.;
  #pragma omp parallel for reduction(+:y[:A->m])
  for (int i=0; i<A->nnz; i++) {
    y[A->row_idx[i]] += A->val[i] * x[A->col_idx[i]];
  }
}

考虑到 COO SpMV 的竞争条件的潜力，这已经是最好的了，而且也没有那么糟糕。实际上，它比 MKL 的

mkl_sparse_d_mv

with COO 做得更好，我认为后者根本没有并行化。

现在，我想并行化稀疏矩阵-矩阵（SpMM）乘积。我做了两次尝试。首先，我写了

void dcoomm_pll1(SparseMatrixCOO *A, double *X, double *Y, int nvec) {
  for (int i=0; i<(A->m * nvec); i++) Y[i] = 0.;
  #pragma omp parallel for reduction(+:Y[:A->m*nvec])
  for (int i=0; i<A->nnz; i++) {
    for (int vec=0; vec<nvec; vec++) {
      Y[A->row_idx[i] + vec * A->m] += A->val[i] * X[A->col_idx[i] + vec * A->n];
    }
  }
}

然而，这并没有加速，反而减慢了并行代码的速度。

第二，我做到了

void dcoomm_pll2(SparseMatrixCOO *A, double *X, double *Y, int nvec) {
  for (int i=0; i<(A->m * nvec); i++) Y[i] = 0.;
  for (int i=0; i<A->nnz; i++) {
    #pragma omp parallel for
    for (int vec=0; vec<nvec; vec++) {
      Y[A->row_idx[i] + vec * A->m] += A->val[i] * X[A->col_idx[i] + vec * A->n];
    }
  }
}

这再次导致速度放缓。

提醒一下，

dcoomm

的顺序实现由

给出

void dcoomm(SparseMatrixCOO *A, double *X, double *Y, int nvec) {
  for (int i=0; i<(A->m * nvec); i++) Y[i] = 0.;
  for (int i=0; i<A->nnz; i++) {
    for (int vec=0; vec<nvec; vec++) {
      Y[A->row_idx[i] + vec * A->m] += A->val[i] * X[A->col_idx[i] + vec * A->n];
    }
  }
}

我的顺序版本比带有 COO 的 MKL

mkl_sparse_d_mm

快约 30%。然而，我的两个多线程版本都比

mkl_sparse_d_mm

慢得多。

当我尝试将COO SpMM函数与

dcoomm_pll1

和

dcoomm_pll2

并行化时，没有一个提供令人满意的性能，我想知道如何并行化COO SpMM以提高并行性能。

Answer 1

矩阵-矩阵乘法代码中的内存访问模式不连续，因此会导致更多缓存未命中。缓存未命中对于顺序代码来说已经很糟糕了，对于多线程代码来说更糟。为什么不调用 nvec 次 dcoomv 的顺序版本，并并行化此循环？你甚至不需要减少。

void dcoomm_pll3(SparseMatrixCOO *A, double *X, double *Y, int nvec) {
  #pragma omp parallel for
  for (int vec=0; vec<nvec; vec++) {
     dcoomv_seq(A, &X[vec * A->n], &Y[vec * A->m])
  }
}