title: Improving parallelization in Numba
body: <p>I have a function in Numba with several loops that could be parallelized. The loop writes to a common array, K, so I understand that the compiler might not be optimizing as much as it could. 
  However, I have no feel for what will make numba's jit compiler create the most efficient code. The examples in the documentation are too simplistic to be of help.</p>
  
  <p>I have tried changing each of the <code>range</code> into <code>prange</code>. The best results were obtained when parallelizing the loop over <code>k_x</code>, but I only got a 2.8x improvement on a machine with 4 cores. I know that I should not expect linear performance improvements, but I have the feeling that I should get better results in this case. For example, I get slightly better results using dask, <code>x_cond.map_blocks(cond_expect_kernel, x_tr, *args)</code> which is odd considering the scheduler overhead.</p>
  
  <p><strong>Is there a way to improve the parallelization of this function beyond simply changing <code>range</code> into <code>prange</code>?</strong></p>
  
  <p>Original function</p>
  
  <pre><code>@jit(float64[:,:](float64[:,:], float64[:,:], int64, int64), nopython=True, nogil=True)
  def cond_expect_kernel(x_cond, x_tr, degree, amount_non_cond_vars):
      size = x_cond.shape[1]
      x_tr_cond = x_tr[:, :size]
      samples_x = x_cond.shape[0]
      samples_tr = x_tr.shape[0]
      K = (1+np.dot(x_cond, x_tr_cond.T))**degree
  
      for j in range(size, size+amount_non_cond_vars):
          for k_x in range(samples_x):
              for k_x_tr in range(samples_tr):
                  K[k_x, k_x_tr] += x_tr[k_x_tr, j]**2*3 
                  for j_left in range(size):
                      K[k_x, k_x_tr] += x_cond[k_x, j_left]*x_tr[k_x_tr, j_left]*x_tr[k_x_tr, j] ** 2 *3  
      return K
  </code></pre>
  
  <p>Best parallel version so far:</p>
  
  <pre><code>@jit(float64[:,:](float64[:,:], float64[:,:], int64, int64), nopython=True, parallel=True)
  def cond_expect_kernel_parallel(x_cond, x_tr, degree, amount_non_cond_vars):
      size = x_cond.shape[1]
      x_tr_cond = x_tr[:, :size]
      samples_x = x_cond.shape[0]
      samples_tr = x_tr.shape[0]
      K = (1+np.dot(x_cond, x_tr_cond.T))**degree
      for j in range(size, size+amount_non_cond_vars):
          for k_x in prange(samples_x):
              for k_x_tr in range(samples_tr):
                  K[k_x, k_x_tr] += x_tr[k_x_tr, j]**2*3 
                  for j_left in range(size):
                      K[k_x, k_x_tr] += x_cond[k_x, j_left]*x_tr[k_x_tr, j_left]*x_tr[k_x_tr, j] ** 2 *3  
      return K
  </code></pre>
  
  <p>For reference, I am working on a machine with 4 cores and another with 16 cores. <code>samples_x</code> is around 100,000, <code>samples_tr</code> around 50000, <code>size</code> around 3, and <code>amount_non_cond_vars</code> around 100.</p>
  
  <p>Thanks!</p>