Skip to content

Commit

Permalink
mpi: Add _allocate_buffers func, drop redundant determinism
Browse files Browse the repository at this point in the history
  • Loading branch information
@georgebisbas
georgebisbas committed Feb 16, 2024
1 parent c85bceb commit ff919e2
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 47 deletions.
67 changes: 22 additions & 45 deletions devito/mpi/routines.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -468,9 +468,11 @@ def _make_haloupdate(self, f, hse, key, sendrecv, **kwargs):
return HaloUpdate('haloupdate%s' % key, iet, parameters)

def _make_basic_mapper(self, f, fixed):
# Build a mapper `(dim, side, region) -> (size, ofs)` for `f`. `size` and
# `ofs` are symbolic objects. This mapper tells what data values should be
# sent (OWNED) or received (HALO) given dimension and side
"""
Build a mapper `(dim, side, region) -> (size, ofs)` for `f`. `size` and
`ofs` are symbolic objects. This mapper tells what data values should be
sent (OWNED) or received (HALO) given dimension and side
"""
mapper = {}
for d0, side, region in product(f.dimensions, (LEFT, RIGHT), (OWNED, HALO)):
if d0 in fixed:
Expand Down Expand Up @@ -636,13 +638,9 @@ def _make_haloupdate(self, f, hse, key, sendrecv, **kwargs):

iet = List(body=body)

parameters = list(f.handles) + [comm, nb] + list(fixed.values())

node = HaloUpdate('haloupdate%s' % key, iet, parameters)

node = node._rebuild(parameters=node.parameters + (kwargs['msg'],))
parameters = list(f.handles) + [comm, nb] + list(fixed.values()) + [kwargs['msg']]

return node
return HaloUpdate('haloupdate%s' % key, iet, parameters)

def _call_haloupdate(self, name, f, hse, msg):
call = super()._call_haloupdate(name, f, hse)
Expand Down Expand Up @@ -1298,6 +1296,17 @@ def _as_number(self, v, args):
assert args is not None
return int(subs_op_args(v, args))

def _allocate_buffers(self, f, shape, entry):
entry.sizes = (c_int*len(shape))(*shape)
size = reduce(mul, shape)*dtype_len(self.target.dtype)
ctype = dtype_to_ctype(f.dtype)
entry.bufg, bufg_memfree_args = self._allocator._alloc_C_libcall(size, ctype)
entry.bufs, bufs_memfree_args = self._allocator._alloc_C_libcall(size, ctype)
# The `memfree_args` will be used to deallocate the buffer upon
# returning from C-land
self._memfree_args.extend([bufg_memfree_args, bufs_memfree_args])
return

def _arg_defaults(self, allocator, alias, args=None):
# Lazy initialization if `allocator` is necessary as the `allocator`
# type isn't really known until an Operator is constructed
Expand All @@ -1315,17 +1324,9 @@ def _arg_defaults(self, allocator, alias, args=None):
except AttributeError:
assert side is CENTER
shape.append(self._as_number(f._size_domain[dim], args))
entry.sizes = (c_int*len(shape))(*shape)

# Allocate the send/recv buffers
size = reduce(mul, shape)*dtype_len(self.target.dtype)
ctype = dtype_to_ctype(f.dtype)
entry.bufg, bufg_memfree_args = allocator._alloc_C_libcall(size, ctype)
entry.bufs, bufs_memfree_args = allocator._alloc_C_libcall(size, ctype)

# The `memfree_args` will be used to deallocate the buffer upon
# returning from C-land
self._memfree_args.extend([bufg_memfree_args, bufs_memfree_args])
self._allocate_buffers(f, shape, entry)

return {self.name: self.value}

Expand Down Expand Up @@ -1376,17 +1377,9 @@ def _arg_defaults(self, allocator, alias, args=None):
except AttributeError:
assert side is CENTER
shape.append(self._as_number(f._size_domain[dim], args))
entry.sizes = (c_int*len(shape))(*shape)

# Allocate the send/recv buffers
size = reduce(mul, shape)*dtype_len(self.target.dtype)
ctype = dtype_to_ctype(f.dtype)
entry.bufg, bufg_memfree_args = allocator._alloc_C_libcall(size, ctype)
entry.bufs, bufs_memfree_args = allocator._alloc_C_libcall(size, ctype)

# The `memfree_args` will be used to deallocate the buffer upon
# returning from C-land
self._memfree_args.extend([bufg_memfree_args, bufs_memfree_args])
self._allocate_buffers(f, shape, entry)

return {self.name: self.value}

Expand Down Expand Up @@ -1444,31 +1437,15 @@ def _arg_defaults(self, allocator, alias, args=None):
# Sending to left, receiving from right
shape = mapper[(d, LEFT, OWNED)]
# Allocate the send/recv buffers
entry.sizes = (c_int*len(shape))(*shape)
size = reduce(mul, shape)*dtype_len(self.target.dtype)
ctype = dtype_to_ctype(f.dtype)
entry.bufg, bufg_memfree_args = allocator._alloc_C_libcall(size, ctype)
entry.bufs, bufs_memfree_args = allocator._alloc_C_libcall(size, ctype)

# The `memfree_args` will be used to deallocate the buffer upon
# returning from C-land
self._memfree_args.extend([bufg_memfree_args, bufs_memfree_args])
self._allocate_buffers(f, shape, entry)

if (d, RIGHT) in self.halos:
entry = self.value[i]
i = i + 1
# Sending to right, receiving from left
shape = mapper[(d, RIGHT, OWNED)]
# Allocate the send/recv buffers
entry.sizes = (c_int*len(shape))(*shape)
size = reduce(mul, shape)*dtype_len(self.target.dtype)
ctype = dtype_to_ctype(f.dtype)
entry.bufg, bufg_memfree_args = allocator._alloc_C_libcall(size, ctype)
entry.bufs, bufs_memfree_args = allocator._alloc_C_libcall(size, ctype)

# The `memfree_args` will be used to deallocate the buffer upon
# returning from C-land
self._memfree_args.extend([bufg_memfree_args, bufs_memfree_args])
self._allocate_buffers(f, shape, entry)

return {self.name: self.value}

Expand Down
2 changes: 0 additions & 2 deletions devito/passes/iet/misc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,8 +140,6 @@ def generate_macros(iet):
# Generate Macros from higher-level SymPy objects
applications = FindApplications().visit(iet)
headers = set().union(*[_generate_macros(i) for i in applications])
# Sort for deterministic code generation
headers = sorted(headers)

# Some special Symbols may represent Macros defined in standard libraries,
# so we need to include the respective includes
Expand Down

0 comments on commit ff919e2

Please sign in to comment.