密集矩阵乘法优化在python中影响很小

Question

I would like to know why the Loop optimization inversion on i and j (which should give better data locality) does not work with the following code in python :

import random
import time

def fillRand1D(size):
    return [random.uniform(-2.0, 2.0) for _ in range(size * size)]

def mmNaive1D(A, B, size):
    C = [0] * size * size
    for i in range(size):
        for j in range(size):
            for k in range(size):
                C[i * size + j] += A[i * size + k] * B[k * size + j]
    return C

def mmInvariant1D(A, B, size):
    C = [0] * size * size
    for i in range(size):
        for j in range(size):
            sigma = C[i * size + j]
            for k in range(size):
                sigma += A[i * size + k] * B[k * size + j]
            C[i * size + j] = sigma
    return C


def mmLoop1D(A, B, size):
    C = [0] * size * size
    for k in range(size):
        for i in range(size):
            for j in range(size):
                C[i * size + j] += A[i * size + k] * B[k * size + j]
    return C


def mmLoopInvariant1D(A, B, size):
    C = [0] * size * size
    for k in range(size):
        for i in range(size):
            Aik = A[i * size + k]
            for j in range(size):
                C[i * size + j] += Aik * B[k * size + j]
    return C

def main():
    matmul_func_1D = [mmNaive1D,mmInvariant1D,mmLoop1D,mmLoopInvariant1D]
    size = 200
    A_1D = fillRand1D(size)
    B_1D = fillRand1D(size)

    for f in matmul_func_1D:
        A = A_1D[:] # copy !
        B = B_1D[:]
        start_time = time.time()
        C = f(A_1D,B_1D,size)
        # print(T)
        print(f.__name__ + " in " + str(time.time() - start_time) + " s")

if __name__ == '__main__':
    main()

Results are with python :

mmNaive1D in 3.420367956161499 s  
mmInvariant1D in 2.316128730773926 s  
mmLoop1D in 3.4071271419525146 s  
mmLoopInvariant1D in 2.5221548080444336 s

Whereas the same optimizations written in C++ give :

> Time [MM naive] 1.780587 s
> Time [MM invariant] 1.642554 s
> Time [MM loop IKJ] 0.304621 s
> Time [MM loop IKJ invariant] 0.276159 s

Answer 1

Your A, B and C matrices are python lists, which are quite slow for big dimensions. You could create them as a numpy arrays to try to speed up the functions:

>>> import numpy as np
>>> A = np.random.uniform(-2., 2., (size,size))
>>> B = np.random.uniform(-2., 2., (size,size))

This defines A and B matrices as size x size random matrices.

The bottleneck of your code are the python loops. Python loops are quite slow, and you want to avoid them in algebraic operations whenever possible. In your case, what you are trying to achieve (if I'm not wrong) is C = A * B where * is the dot product between two 2D matrices. Translated to numpy:

>>> C = A.dot(B)

If you time it:

>>> %timeit C = A.dot(B)
100 loops, best of 3: 7.91 ms per loop

Takes only 7.91 ms to calculate the dot product. I know that you wanted to play with python, but if you are going to do mathematical calculus, the sooner you move to numpy the better for your algorithms' speeds.

Going back to your algorithm, if I run your code in my computer, the optimization does work and achieves a bit of speedup:

mmNaive1D in 1.72708702087 s
mmInvariant1D in 1.64227509499 s
mmLoop1D in 1.57529997826 s
mmLoopInvariant1D in 1.26218104362 s

So it is working for me.

EDIT: After runing the code several times I'm getting different results:

mmNaive1D in 1.63492894173 s
mmInvariant1D in 1.1577808857 s
mmLoop1D in 1.67409181595 s
mmLoopInvariant1D in 1.32283711433 s

This time the "optimization" is not working. And I guess it is because the bottleneck of the algorithm are the python loops, not the memory accesses.

EDIT2: Some more relevant results performing 100 iterations of each algorithm to get mean timings (took it a while haha):

mmNaive1D in 1.66692941904 s
mmInvariant1D in 1.15141540051 s
mmLoop1D in 1.58852998018 s
mmLoopInvariant1D in 1.28386260986 s

Even if 100 iterations are still not revelant enough (like 1.000.000 would be better, but it would take ages), it suggest that as the mean of 100 runtimes, the "optimization" is not really an optimization. I don't really know why, but it could be that adding a new variable in python is slower than 200 memory accesses (maybe there is some cache behind it).