Thursday, August 25, 2011

Concurrent matrix multiplication in Python

We're going to demonstrate parallel matrix multiplication in Python.

Lets suppose we are performing the multiplication: P = A * B. We're going to calculate the result matrix row-by-row. The following function returns a row of a result matrix P. The arguments are a row of matrix A and the matrix B.

def calc_row_of_product_matrix(a_row, b):
    return map(lambda col: sum(starmap(mul,izip(a_row,col))), izip(*b))


We're going to use the multiprocessing module of Python to create a process pool containing as many processes as the number of CPUs we have and calculate each result row in a different process. Here are the two lines that does this trick:

def __mul__(self, b):
    pool = multiprocessing.Pool(multiprocessing.cpu_count())
    return pool.map(eval_func_tuple, izip(repeat(calc_row_of_product_matrix), self, repeat(b)))


Some explanation why is this code a bit nasty. Functions passed to the pool has to be picklable and because only the functions defined at the top level of a module are picklable, we needed to get the the row calculation function calc_row_of_product_matrix out of the class. Furthermore, to be able to pass two arguments to calc_row_of_product_matrix we also needed a helper function, which takes a tuple of a function and args, evaluates and returns the result:

def eval_func_tuple(f_args):
    return f_args[0](*f_args[1:])


Note that passing around the whole matrix B to all processes calculating result rows, or more precisely passing itertools.repeat(b) to  pool.map() visibly increases the memory consumption of the multiprocess version. This was not a real issue for me, as the bottleneck was CPU; anyway, this issue could be addressed by using the shared memory module of multiprocessing. For now we'll leave that as an exercise to the reader.

Here are the running times on my Intel Core2 Duo 3.16GHz box. For sufficiently large matrix (above 500*500) the running time of the multiprocess version is nearly half of the single process version.


100*100 matrix, single process: 0.0835670982343 
100*100 matrix, multiprocess: 0.351096555199 
200*200 matrix, single process: 0.79961114284 
200*200 matrix, multiprocess: 0.700980680681 
500*500 matrix, single process: 14.4003259234 
500*500 matrix, multiprocess: 7.99582457187 
1000*1000 matrix, single process: 118.078526896 
1000*1000 matrix, multiprocess: 66.8809939919 

Here you can see the single process version running with CPU usage of 50%:

Here you can see  the multiprocessing version running with CPU usage of 100%:



Here's the full code: