Matrix Multiplication Implementation

[qedrohenrique]

I'm working in implementing a matrix multiplication with Charm4Py.
Currently, i'm using Pools to do it, but it isn't that fast. I'm wondering how to improve it.
Any ideas?

from charm4py import charm, coro
import numpy as np
import time


MATRIX_SIZE = 240


@coro
def matrix_multiply(data):
    ma, mb, i, j = data
    result = 0
    for k in range(MATRIX_SIZE):
        result += ma[i][k] * mb[k][j]
    return result


def main(args):
    ma = np.random.randint(10, size=(MATRIX_SIZE, MATRIX_SIZE))
    mb = np.random.randint(10, size=(MATRIX_SIZE, MATRIX_SIZE))

    tasks = [(ma, mb, i, j) for i in range(MATRIX_SIZE) for j in range(MATRIX_SIZE)]

    start_time = time.time()
    result = charm.pool.map(
        matrix_multiply,
        tasks
    )
    end_time = (time.time() - start_time)

    print("[%dx%d]" % (MATRIX_SIZE, MATRIX_SIZE))
    print("%d processes" % charm.numPes())
    print("Elapsed time: %.4f seconds" % end_time)
    print("%.4f\n" % end_time)

    # matrix_result = np.array(result).reshape(MATRIX_SIZE, MATRIX_SIZE)
    # print(matrix_result)
    exit(0)


charm.start(main)

For reference, using Intel(R) Core(TM) i7-10700KF CPU @ 3.80GHz with 16 PEs.

[240x240]
16 processes
Elapsed time: 42.7994 seconds
42.7994

[ZwFink]

Hello,
Any advice on improving your implementation depends on what you're trying to do.
Do you want a faster, distributed matrix multiply? If so, you can move to an algorithm like Cannon/SUMMA.
Do you want to improve the speed of the above code? If so, you can use something like NUMBA for the dot product, replace your loop with some vectorized numpy version, increase the granularity of the tasks, etc. Some measurements of where time is spent would help here, but I suspect a large problem is creating $MATRIXSIZE^2$ tasks is harming performance.

[qedrohenrique]

Thanks for the advice. Since i'm not working with a distributed matrix, Cannon's algorithm wasn't the thing i was looking for.
But, on the other hand, i've managed to reach an ok result using JIT compilation and Chares.

This is what I got so far.
(I know there's improvements to be done, I'll get there)

class MatrixMultiply(Chare):

    def __init__(self, ma, mb, sim_done_future):
        dimension = MATRIX_SIZE // charm.numPes()

        self.ma = ma
        self.mb = mb
        self.mc = np.zeros(shape=(dimension, MATRIX_SIZE))
        self.sim_done_future = sim_done_future

        self.start = charm.myPe() * dimension
        self.end = self.start + dimension - 1


    @coro
    def run(self):
        # print('PE [%d] - %d to %d' % (charm.myPe(), self.start, self.end))

        self.mc = compute(self.ma, self.mb, self.mc, self.start, self.end)

        self.reduce(self.sim_done_future, self.mc.tolist(), Reducer.gather)  # Sync all Chares


@jit(nopython=True, cache=False)
def compute(ma, mb, mc, start, end):
    for i in range(start, end + 1):
        for j in range(MATRIX_SIZE):
            for k in range(MATRIX_SIZE):
                mc_i = i - start
                mc[mc_i][j] += ma[i][k] * mb[k][j]

    return mc

[2048x2048]
1 processes
Elapsed time: 31.2518 seconds
31.2518