Efficient way to calculate 3D matrix multiplication using numpy - numpy

Question

How can I efficiently write and calculate this multiplication using numpy:
for k in range(K):
for i in range(SIZE):
for j in range(SIZE):
for i_b in range(B_SIZE):
for j_b in range(B_SIZE):
for k_b in range(k+1):
data[k, i * w + i_b, j * h + j_b] += arr1[k_b, i_b, j_b] * arr2[k_b, i, j]
For example:
SIZE, B_SIZE = 32, 8
arr1.shape -> (8, 8, 8)
arr2.shape -> (8, 32, 32)
data.shape -> (K, 256, 256)
Thank you.

Answer 1

You can use Numba for such kind of non-trivial case and rework the loops to use efficiently the CPU cache. Here is an example:
import numba as nb
#nb.njit
def compute(data, arr1, arr2):
for k in range(K):
for k_b in range(k+1):
for i in range(SIZE):
for j in range(SIZE):
tmp = arr2[k_b, i, j]
for i_b in range(B_SIZE):
for j_b in range(B_SIZE):
data[k, i * w + i_b, j * h + j_b] += arr1[k_b, i_b, j_b] * tmp
If you do this operation once, then you can pre-compile the Numba code by providing the types of the arrays. If K is big, then you can parallelize the code using #nb.njit(parallel=True) and use for k in nb.prange(K) rather than for k in range(K). This should be several order of magnitude fater.