title: How to parallelize this Python for loop when using Numba
body: <p>I'm using the Anaconda distribution of Python, together with Numba, and I've written the following Python function that multiplies a sparse matrix <strong><code>A</code></strong> (stored in a CSR format) by a dense vector <strong><code>x</code></strong>:</p>
  
  <pre><code>@jit
  def csrMult( x, Adata, Aindices, Aindptr, Ashape ):
  
      numRowsA = Ashape[0]
      Ax       = numpy.zeros( numRowsA )
  
      for i in range( numRowsA ):
          Ax_i = 0.0
          for dataIdx in range( Aindptr[i], Aindptr[i+1] ):
  
              j     = Aindices[dataIdx]
              Ax_i +=    Adata[dataIdx] * x[j]
  
          Ax[i] = Ax_i
  
      return Ax 
  </code></pre>
  
  <p>Here <strong><code>A</code></strong> is a large <code>scipy</code> sparse matrix,</p>
  
  <pre><code>&gt;&gt;&gt; A.shape
  ( 56469, 39279 )
  #                  having ~ 142,258,302 nonzero entries (so about 6.4% )
  &gt;&gt;&gt; type( A[0,0] )
  dtype( 'float32' )
  </code></pre>
  
  <p>and <strong><code>x</code></strong> is a <code>numpy</code> array. Here is a snippet of code that calls the above function:</p>
  
  <pre><code>x       = numpy.random.randn( A.shape[1] )
  Ax      = A.dot( x )   
  AxCheck = csrMult( x, A.data, A.indices, A.indptr, A.shape )
  </code></pre>
  
  <p>Notice the <strong><code>@jit</code></strong>-decorator that tells Numba to do a just-in-time compilation for the <strong><code>csrMult()</code></strong> function.</p>
  
  <p>In my experiments, my function <code>csrMult()</code> is about <strong>twice as fast</strong> as the <code>scipy</code> <strong><code>.dot()</code></strong> method.  That is a pretty impressive result for Numba.</p>
  
  <p>However, MATLAB still performs this matrix-vector multiplication about <strong>6 times faster</strong> than <code>csrMult()</code>.  I believe that is because MATLAB uses multithreading when performing sparse matrix-vector multiplication.</p>
  
  <hr>
  
  <h2><strong>Question:</strong></h2>
  
  <p>How can I parallelize the outer <code>for</code>-loop when using Numba?</p>
  
  <p>Numba used to have a <strong><code>prange()</code></strong> function, that made it simple to parallelize embarassingly parallel <strong><code>for</code></strong>-loops. Unfortunately, Numba no longer has <code>prange()</code> [<strong>actually, that is false, see the edit below</strong>].  <strong>So what is the correct way to parallelize this <code>for</code>-loop now, that Numba's <code>prange()</code> function is gone?</strong></p>
  
  <p>When <code>prange()</code> was removed from Numba, what alternative did the developers of Numba have in mind?</p>
  
  <hr>
  
  <blockquote>
    <p><strong>Edit 1:</strong><br>
    I updated to the latest version of Numba, which is .35, and <code>prange()</code> is back! It was not included in version .33, the version I had been using.<br>
    That is good news, but unfortunately I am getting an error message when I attempt to parallelize my for loop using <code>prange()</code>.  Here is a parallel for loop <a href="http://numba.pydata.org/numba-doc/latest/user/parallel.html" rel="noreferrer">example</a> from the Numba documentation (see section 1.9.2 "Explicit Parallel Loops"), and below is my new code:</p>
  </blockquote>
  
  <pre><code>from numba import njit, prange
  @njit( parallel=True )
  def csrMult_numba( x, Adata, Aindices, Aindptr, Ashape):
  
      numRowsA = Ashape[0]    
      Ax       = np.zeros( numRowsA )
  
      for i in prange( numRowsA ):
          Ax_i = 0.0        
          for dataIdx in range( Aindptr[i],Aindptr[i+1] ):
  
              j     = Aindices[dataIdx]
              Ax_i +=    Adata[dataIdx] * x[j]
  
          Ax[i] = Ax_i            
  
      return Ax 
  </code></pre>
  
  <p>When I call this function, using the code snippet given above, I receive the following error: </p>
  
  <blockquote>
    <p>AttributeError: Failed at nopython (convert to parfors) 'SetItem'
    object has no attribute 'get_targets'</p>
  </blockquote>
  
  <hr>
  
  <h2>Given<br>the above attempt to use <code>prange</code> crashes, my question stands:</h2>
  
  <p><strong>What is the correct way</strong> ( using <code>prange</code> or an alternative method ) <strong>to parallelize this Python <code>for</code>-loop?</strong>  </p>
  
  <p>As noted below, it was trivial to parallelize a similar for loop in C++ and obtain an <strong>8x</strong> speedup, having been run on <strong>20</strong>-omp-threads.  There must be a way to do it using Numba, since the for loop is embarrassingly parallel (and since sparse matrix-vector multiplication is a fundamental operation in scientific computing). </p>
  
  <hr>
  
  <blockquote>
    <p><strong>Edit 2:</strong><br>
    Here is my C++ version of <code>csrMult()</code>.  Parallelizing the <code>for()</code> loop in the C++ version makes the code about 8x faster in my tests. This suggests to me that a similar speedup should be possible for the Python version when using Numba.</p>
  </blockquote>
  
  <pre><code>void csrMult(VectorXd&amp; Ax, VectorXd&amp; x, vector&lt;double&gt;&amp; Adata, vector&lt;int&gt;&amp; Aindices, vector&lt;int&gt;&amp; Aindptr)
  {
      // This code assumes that the size of Ax is numRowsA.
      #pragma omp parallel num_threads(20)
      {       
          #pragma omp for schedule(dynamic,590) 
          for (int i = 0; i &lt; Ax.size(); i++)
          {
              double Ax_i = 0.0;
              for (int dataIdx = Aindptr[i]; dataIdx &lt; Aindptr[i + 1]; dataIdx++)
              {
                  Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
              }
  
              Ax[i] = Ax_i;
          }
      }
  }
  </code></pre>