Importing pykokkos serializes scipy ffts on AMD EPYC 7763 64-Core Processor (i.e., TACC's ls6)

[kennykos]

Using the following test script

from time import time
import numpy as np
from scipy.fft import rfftn
import pykokkos

grid = np.random.normal(0,1, (96, 116, 116, 12))
x = np.random.normal(0, 1, (96, 292, 292, 12))
s= (96, 292, 292)
tic = time()                                                                                                                                                                
xfft = rfftn(grid.transpose(3,0,1,2), axes = (1, 2, 3), s=s, workers = 128)
toc = time()
print('time warm up: ', toc - tic)
 
repeat = 2
 
tic_all = time()
for i in range(repeat):
    tic = time()                                                                                                                                                                 
    xfft = rfftn(grid.transpose(3,0,1,2), axes = (1, 2, 3), overwrite_x=True, s=s, workers = 128)
    toc = time()
    print('time: ', toc - tic)
toc_all = time()
 
print('time ave: ', (toc_all - tic_all)/repeat)

when I comment out import pykokkos I get

time warm up:  0.16707849502563477
time:  0.15308403968811035
time ave:  0.07656288146972656
time:  0.14578652381896973
time ave:  0.14947926998138428

but when I leave import pykokkos uncommented I get

time warm up:  1.8171048164367676
time:  1.8107686042785645
time ave:  0.9054266214370728
time:  1.8127191066741943
time ave:  1.8118274211883545

[IvanGrigorik]

Is it repetitive program runtime or time of first run?

[kennykos]

Is it repetitive program runtime or time of first run?

The times reported above are for 2 repeats, but this is an instructive behavior of what we observe during a more complex program where an FFT is called inside a loop. I discussed with @wlruys, and I'm going to check and see if my pykokkos install is dynamically linked to any of the same libraries as scpipy compiled with different flags.

I don't see these results on Vista which has a Grace CPU, so it may be an x86 issue (just a guess)?

[kennykos]

If I change the script to

from time import time
import numpy as np
from scipy.fft import rfftn
rfftn(np.ones((2,2,2,2)), workers=128) # run an rfftn before importing pykokkos
import pykokkos

grid = np.random.normal(0,1, (96, 116, 116, 12))
x = np.random.normal(0, 1, (96, 292, 292, 12))
s= (96, 292, 292)
tic = time()
xfft = rfftn(np.ones((2,2,2,2)), axes = (1, 2, 3), s=s, workers = 128)
toc = time()
print('time warm up: ', toc - tic)

repeat = 2

tic_all = time()
for i in range(repeat):
  tic = time()
  xfft = rfftn(grid.transpose(3,0,1,2), axes = (1, 2, 3), overwrite_x=True, s=s, workers = 128)
  toc = time()
  print('time: ', toc - tic)
  toc_all = time()
  print('time ave: ', (toc_all - tic_all)/repeat)

my times drop to

time warm up:  0.026160717010498047
time:  0.15665459632873535
time ave:  0.07834994792938232
time:  0.15795397758483887
time ave:  0.1573486328125

an 11.45x speedup

[JBludau]

Are you using openmp? it might be that we set some variables on the pool if they are not set before. Scipy might do that on first run of something that uses the pool

[kennykos]

Are you using openmp? it might be that we set some variables on the pool if they are not set before. Scipy might do that on first run of something that uses the pool

Yes, I'm using OpenMP. How can I check what pool variables are being set? Also, If I change the import pykokkos line to import kokkos in my first script (without running rfftn before importing) I get the same 11x speedup as in my second script.