Broadcasting outside main loop speeds up vectorized numpy ops?

Question

I'm doing some vectorized algebra using numpy and the wall-clock performance of my algorithm seems weird. The program does roughly as follows:

1. Create three matrices: Y (KxD), X (NxD), T (KxN)
2. For each row of Y:
3. subtract Y[i] from each row of X (by broadcasting),
4. square the differences along one axis, sum them, take a square root, then store in T.

However, depending on how I perform the broadcasting, computation speed is vastly different. Consider the code:

import numpy as np
from time import perf_counter

D = 128
N = 3000
K = 500

X = np.random.rand(N, D)
Y = np.random.rand(K, D)
T = np.zeros((K, N))

if True: # negate to enable the second loop
    time = 0.0
    for i in range(100):
        start = perf_counter()
        for i in range(K):
            T[i] = np.sqrt(np.sum(
                np.square(
                  X - Y[i] # this has dimensions NxD
                ),
                axis=1
            ))
        time += perf_counter() - start
    print("Broadcast in line: {:.3f} s".format(time / 100))
    exit()

if True:
    time = 0.0
    for i in range(100):
        start = perf_counter()
        for i in range(K):
            diff = X - Y[i]
            T[i] = np.sqrt(np.sum(
                np.square(
                  diff
                ),
                axis=1
            ))
        time += perf_counter() - start
    print("Broadcast out:     {:.3f} s".format(time / 100))
    exit()

Times for each loop are measured individually and averaged over 100 executions. The results:

Broadcast in line: 1.504 s
Broadcast out:     0.438 s

The only difference is that broadcasting and subtraction in the first loop is done in-line, while in the second approach I do it before any vectorized operations. Why is this making such a difference?

My system configuration:

• Lenovo ThinkStation P920, 2x Xeon Silver 4110, 64 GB RAM
• Xubuntu 18.04.2 LTS (bionic)
• Python 3.7.3 (GCC 7.3.0)
• Numpy 1.16.3 linked against OpenBLAS (that's as much as np.__config__.show() tells me)

---

PS: Yes I am aware this could be further optimized, but right now I would like to understand what happens under the hood here.

Answer 1

It's not a broadcasting problem

I also added a optimized solution to see how long the actual calculation takes without the large overhead of memory allocation and deallocation.

Functions

import numpy as np
import numba as nb

def func_1(X,Y,T):
    for i in range(K):
        T[i] = np.sqrt(np.sum(np.square(X - Y[i]),axis=1))
    return T

def func_2(X,Y,T):
    for i in range(K):
        diff = X - Y[i]
        T[i] = np.sqrt(np.sum(np.square(diff),axis=1))
    return T

@nb.njit(fastmath=True,parallel=True)
def func_3(X,Y,T):
    for i in nb.prange(Y.shape[0]):
        for j in range(X.shape[0]):
            diff_sq_sum=0.
            for k in range(X.shape[1]):
                diff_sq_sum+= (X[j,k] - Y[i,k])**2
            T[i,j]=np.sqrt(diff_sq_sum)
    return T

Timings

I did all the timings in a Jupyter Notebook and observed a really weird behavior. The following code is in one cell. I also tried calling timit multiple times, but on the first execution of the cell this doesn't change anything.

First execution of the cell

D = 128
N = 3000
K = 500

X = np.random.rand(N, D)
Y = np.random.rand(K, D)
T = np.zeros((K, N))

#You can do it more often it would not change anything
%timeit func_1(X,Y,T)
%timeit func_1(X,Y,T)

#You can do it more often it would not change anything
%timeit func_2(X,Y,T)
%timeit func_2(X,Y,T)

###Avoid measuring compilation overhead###
%timeit func_3(X,Y,T)
##########################################
%timeit func_3(X,Y,T)

774 ms ± 6.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
768 ms ± 2.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
494 ms ± 2.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
494 ms ± 1.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
10.7 ms ± 1.25 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
6.74 ms ± 39.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Second execution

345 ms ± 16.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
337 ms ± 3.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
322 ms ± 834 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
323 ms ± 1.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
6.93 ms ± 234 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
6.9 ms ± 87.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)