我的 CSR 稀疏矩阵多向量 (SpMM) 乘积函数有什么问题？

Question

我有以下 C 语言稀疏矩阵向量 (SpMV) 乘积代码（假设采用 CSR 存储格式）：

void dcsrmv(SparseMatrixCSR *A, double *x, double *y) {
  for (int i=0; i<A->m; i++) {
    double sum = 0.;
    for (int jj=A->row_start[i]; jj<A->row_start[i+1]; jj++) {
      sum += A->val[jj] * x[A->col_idx[jj]];
    }
    y[i] = sum;
  }
}

当我使用 SuiteSparse Matrix Collection 中的数组

ecology1

运行此函数 100 次时，平均运行时间为 5.9 毫秒，而如果我使用相同的矩阵运行 MKL 的

mkl_sparse_d_mv

函数，则平均运行时间为 8.4 毫秒。

现在，我想要一个具有相同存储格式的稀疏矩阵多向量 (SpMM) 乘积函数，我对其进行了以下 3 种方法的编码：

void dcsrmm0(SparseMatrixCSR *A, double *X, double *Y, int nvec) {
  for (int vec=0; vec<nvec; vec++) {
    double *x = X + vec * A->n;
    double *y = Y + vec * A->m;
    for (int i=0; i<A->m; i++) {
      double sum = 0.;
      for (int jj=A->row_start[i]; jj<A->row_start[i+1]; jj++) {
        sum += A->val[jj] * x[A->col_idx[jj]];
      }
      y[i] = sum;
    }
  }
}

void dcsrmm1(SparseMatrixCSR *A, double *X, double *Y, int nvec) {
  for (int i=0; i<A->m; i++) {
    for (int vec=0; vec<nvec; vec++) {
      double sum = 0.;
      double *x = X + vec * A->n;
      double *y = Y + vec * A->m;
      for (int jj=A->row_start[i]; jj<A->row_start[i+1]; jj++) {
        sum += A->val[jj] * x[A->col_idx[jj]];
      }
      y[i] = sum;
    }
  }
}

和

void dcsrmm2(SparseMatrixCSR *A, double *X, double *Y, int nvec) {
  double *x[nvec], *y[nvec];
  for (int vec=0; vec<nvec; vec++) x[vec] = X + vec * A->n;
  for (int vec=0; vec<nvec; vec++) y[vec] = Y + vec * A->m;
  double sum[nvec];
  for (int i=0; i<A->m; i++) {
    for (int vec=0; vec<nvec; vec++) sum[vec] = 0.;
    for (int jj=A->row_start[i]; jj<A->row_start[i+1]; jj++) {
      for (int vec=0; vec<nvec; vec++) {
        sum[vec] += A->val[jj] * x[vec][A->col_idx[jj]];
      }
    }
    for (int vec=0; vec<nvec; vec++) y[vec][i] = sum[vec];
  }
}

其中

X

和

Y

按列存储。

然后我用

nvec=10

对这些函数进行基准测试。在这种情况下，

dcsrmm0

平均需要 49.8 毫秒，

dcsrmm1

平均需要 38.3 毫秒，

dcsrmm2

平均需要 34.8 毫秒，而

mkl_sparse_d_mm

仅需要 26.7 毫秒。

所以我的问题是，为什么我的

dcsrmv

实现比

mkl_sparse_d_mv

更好，但我的

dcsrmm0-2

功能却落后于

mkl_sparse_d_mm

那么多？我可以做什么来改进这些功能？

Answer 1

我发现以下替代方案给了我最好的结果：

void dcsrmm2_improved(SparseMatrixCSR *A, double *X, double *Y, int nvec) {
  double *x[nvec];
  for (int vec=0; vec<nvec; vec++) x[vec] = X + vec * A->n;
  for (int i=0; i<A->m; i++) {
    double *y = Y + i;
    for (int vec=0; vec<nvec; vec++) y[vec * A->m] = 0.;
    for (int jj=A->row_start[i]; jj<A->row_start[i+1]; jj++) {
      double val = A->val[jj];
      int j = A->col_idx[jj];
      for (int vec=0; vec<nvec; vec++) {
        y[vec * A->m] += val * x[vec][j];
      }
    }
  }
}

当我使用 SuiteSparse Matrix Collection 中的数组

ecology1

运行此函数 100 次时，我得到的平均运行时间为 27.6 毫秒，这相当接近使用 CSR 格式使用

mkl_sparse_d_mm

获得的 26.7 毫秒。

如果有人知道如何进一步改进，请告诉我。