Estoy usando la distribución Anaconda de Python, junto con Numba, y he escrito la siguiente función de Python que multiplica una matriz dispersaA (almacenado en formato CSR) por un vector densox:

@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 

aquíA es grandescipy matriz dispersa,

>>> A.shape
( 56469, 39279 )
#                  having ~ 142,258,302 nonzero entries (so about 6.4% )
>>> type( A[0,0] )
dtype( 'float32' )

yx es unnumpy formación. Aquí hay un fragmento de código que llama a la función anterior:

x       = numpy.random.randn( A.shape[1] )
Ax      = A.dot( x )   
AxCheck = csrMult( x, A.data, A.indices, A.indptr, A.shape )

Observe la@jit-decorador que le dice a Numba que haga una compilación justo a tiempo para elcsrMult() función.

En mis experimentos, mi funcióncsrMult() es sobreel doble de rápido como elscipy .dot() método. Ese es un resultado bastante impresionante para Numba.

Sin embargo, MATLAB todavía realiza esta multiplicación de matriz-vector sobre6 veces más rápido quecsrMult(). Creo que esto se debe a que MATLAB utiliza subprocesos múltiples al realizar una multiplicación de matriz-vector dispersa.

Pregunta:

¿Cómo puedo paralelizar el exteriorfor-loop cuando se usa Numba?

Numba solía tener unprange() función, que simplificaba la paralelización embarazosamente paralelaforbucles Desafortunadamente, Numba ya no tieneprange() [en realidad, eso es falso, mira la edición a continuación]Entonces, ¿cuál es la forma correcta de paralelizar esto?for-lazo ahora, ese de Numbaprange() la función se ha ido?

Cuandoprange() fue eliminado de Numba, ¿qué alternativa tenían en mente los desarrolladores de Numba?

Editar 1:
Actualicé a la última versión de Numba, que es .35, yprange() ¡está de vuelta! No estaba incluido en la versión .33, la versión que había estado usando.
Es una buena noticia, pero desafortunadamente recibo un mensaje de error cuando intento paralelizar mi bucle for usandoprange(). Aquí hay un paralelo para el bucleejemplo de la documentación de Numba (consulte la sección 1.9.2 "Bucles paralelos explícitos"), y debajo está mi nuevo código:

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 

Cuando llamo a esta función, usando el fragmento de código proporcionado anteriormente, recibo el siguiente error:

AttributeError: Error en nopython (convertir a parfors) El objeto 'SetItem' no tiene el atributo 'get_targets'

Dado
el intento anterior de usarprange se bloquea, mi pregunta es:

Cuál es la manera correcta ( utilizandoprange o un método alternativo)para paralelizar este Pythonfor-¿lazo?

Como se observa a continuación, fue trivial paralelizar un bucle similar en C ++ y obtener un8x acelerar, habiendo sido ejecutado en20-omp-hilos. Debe haber una manera de hacerlo usando Numba, ya que el bucle for es vergonzosamente paralelo (y dado que la multiplicación dispersa de vectores de matriz es una operación fundamental en la computación científica).

Edición 2:
Aquí está mi versión C ++ decsrMult(). Paralelizando elfor() Loop en la versión C ++ hace que el código sea aproximadamente 8 veces más rápido en mis pruebas. Esto me sugiere que debería ser posible una aceleración similar para la versión de Python cuando se usa Numba.

void csrMult(VectorXd& Ax, VectorXd& x, vector<double>& Adata, vector<int>& Aindices, vector<int>& 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 < 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;
        }
    }
}