How can I speed up this code for sparse matrix-vector multiplication?

Question

I've written a C++ function that multiplies a sparse matrix (stored in CSR format) by a dense vector. Here's the code:

VectorXd csrMult(VectorXd x, vector<double> Adata, vector<int> Aindices, vector<int> Aindptr, int numRowsA) 
{
    VectorXd Ax = VectorXd::Zero(numRowsA); 
    for (int i = 0; i < numRowsA; i++)
    {
        for (int dataIdx = Aindptr[i]; dataIdx < Aindptr[i + 1]; dataIdx++)
        {
            Ax[i] += Adata[dataIdx] * x[Aindices[dataIdx]];
        }
    }

    return Ax;
}

Here VectorXd is a data type provided by the Eigen3 linear algebra library. The inputs Adata, Aindices, and Aindptr describe a matrix A in CSR format.

To be precise: Adata is a list of nonzero entries of A (stored in row major order), Aindices[i] tells us which column of A the nonzero entry Adata[i] belongs to, and Aindptr[i+1] - Aindptr[i] is the number of nonzero entries in the ith row of A.

I'm observing that Eigen3's sparse matrix-vector multiplication operation is about 5 times faster than my csrMult function, even when openMP is disabled. However, when I look at the source code SparseDenseProduct.h that I believe Eigen3 is using to compute this matrix-vector product, it's not clear to me how Eigen3 is faster.

Question: Do you have any suggestions to improve the speed of my code? Can you explain why Eigen3 is faster?

---

Edit 3: I noticed an important clue about what's going on. When I make the problem smaller by setting all but the top N entries of A equal to 0 (so that A becomes more sparse), my implementation compares better with Eigen3. In fact, if I set all but the top 10,000 nonzero entries of A equal to 0, my implementation beats Eigen3 by a factor of about 3. However, for the very large problem size I am interested in, Eigen3 beats my implementation by a factor of 2. I would still really like to tie with Eigen3 for large problem sizes.

---

Edit 2: I changed the inputs to pass by reference, as @TylerOlsen suggested, and now Eigen3 is only about twice as fast as my code (with openMP disabled). So that was a significant improvement. Here's the latest version of my code. I still need to figure out how to make my code twice as fast in order to tie with Eigen3.

void csrMult_v3(VectorXd& Ax, VectorXd& x, vector<double>& Adata, vector<int>& Aindices, vector<int>& Aindptr)
{
    // This code assumes that the size of Ax is numRowsA.
    for (int i = 0; i < Ax.size(); i++)
    {
        double Ax_i = 0.0;
        for (int dataIdx = Aindptr[i]; dataIdx < Aindptr[i + 1]; dataIdx++)
        {
            Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
        }

        Ax[i] = Ax_i;
    }       
}

---

Edit 1: I also tried the following code, where the value Ax[i] is accumulated in a temporary variable, but the effect on runtime was negligible. I'm still observing Eigen3 is about 5 times faster (without openMP enabled).

VectorXd csrMult_v2(VectorXd x, vector<double> Adata, vector<int> Aindices, vector<int> Aindptr, int numRowsA)
{
    VectorXd Ax = VectorXd::Zero(numRowsA);
    for (int i = 0; i < numRowsA; i++)
    {
        double Ax_i = 0.0;
        for (int dataIdx = Aindptr[i]; dataIdx < Aindptr[i + 1]; dataIdx++)
        {
            Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
        }

        Ax[i] = Ax_i;
    }

    return Ax;
}

Answer 1

After your second edit, you are already here:

for (int i = 0; i < Ax.size(); i++)
{
  double Ax_i = 0.0;
  for (int dataIdx = Aindptr[i]; dataIdx < Aindptr[i + 1]; dataIdx++)
  {
     Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
  }

  Ax[i] = Ax_i;
}

But notice how you in row i+1 you are initializing dataIdx to the value Aindptr[i+1] it already has because you stopped the loop in row i when it had that value. So you can transform this to the following:

int dataIdx = Aindptr[0];
for (int i = 0; i < Ax.size(); i++)
{
  double Ax_i = 0.0;
  for (; dataIdx < Aindptr[i + 1]; dataIdx++)
  {
     Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
  }

  Ax[i] = Ax_i;
}

The next piece of information you need to know is that you are walking the Aindices array linearly, from left to right because your indices dataIdx are incremented one by one. So you can save yourself the subscripting and do this:

int dataIdx = Aindptr[0];
int *Aindex = &Aindices[dataIdx];
for (int i = 0; i < Ax.size(); i++)
{
  double Ax_i = 0.0;
  for (; dataIdx < Aindptr[i + 1]; ++dataIdx,++Aindex)
  {
     Ax_i += Adata[dataIdx] * x[*Aindex];
  }

  Ax[i] = Ax_i;
}

Give this a try.