Skip to content

Commit

Permalink
Added cache for dealiasing to avoid communication
Browse files Browse the repository at this point in the history
  • Loading branch information
brownbaerchen committed Oct 28, 2024
1 parent 2f84eda commit 93e5ac9
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions pySDC/helpers/spectral_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -882,6 +882,7 @@ def __init__(self, comm=None, useGPU=False, debug=False):
self.BCs = None

self.fft_cache = {}
self.fft_dealias_shape_cache = {}

@property
def u_init(self):
Expand Down Expand Up @@ -1470,7 +1471,9 @@ def _transform_dct(self, u, axes, padding=None, **kwargs):

if padding is not None:
shape = list(v.shape)
if self.comm:
if ('forward', *padding) in self.fft_dealias_shape_cache.keys():
shape[0] = self.fft_dealias_shape_cache[('forward', *padding)]
elif self.comm:
send_buf = np.array(v.shape[0])
recv_buf = np.array(v.shape[0])
self.comm.Allreduce(send_buf, recv_buf)
Expand Down Expand Up @@ -1645,7 +1648,9 @@ def _transform_idct(self, u, axes, padding=None, **kwargs):
if padding is not None:
if padding[axis] != 1:
shape = list(v.shape)
if self.comm:
if ('backward', *padding) in self.fft_dealias_shape_cache.keys():
shape[0] = self.fft_dealias_shape_cache[('backward', *padding)]
elif self.comm:
send_buf = np.array(v.shape[0])
recv_buf = np.array(v.shape[0])
self.comm.Allreduce(send_buf, recv_buf)
Expand Down

0 comments on commit 93e5ac9

Please sign in to comment.