def _make_poke(self, hs, key, msgs):
lflag = Symbol(name='lflag')
gflag = Symbol(name='gflag')
# Init flags
body = [Expression(DummyEq(lflag, 0)), Expression(DummyEq(gflag, 1))]
# For each msg, build an Iteration calling MPI_Test on all peers
for msg in msgs:
dim = Dimension(name='i')
msgi = IndexedPointer(msg, dim)
rrecv = Byref(FieldFromComposite(msg._C_field_rrecv, msgi))
testrecv = Call(
'MPI_Test',
[rrecv, Byref(lflag),
Macro('MPI_STATUS_IGNORE')])
rsend = Byref(FieldFromComposite(msg._C_field_rsend, msgi))
testsend = Call(
'MPI_Test',
[rsend, Byref(lflag),
Macro('MPI_STATUS_IGNORE')])
update = AugmentedExpression(DummyEq(gflag, lflag), '&')
body.append(
Iteration([testsend, update, testrecv, update], dim,
msg.npeers - 1))
body.append(Return(gflag))
return make_efunc('pokempi%d' % key, List(body=body), retval='int')
def _make_halowait(self, f, hse, key, msg=None):
nb = f.grid.distributor._obj_neighborhood
wait = self._cache_dims[f.dimensions][2]
fixed = {d: Symbol(name="o%s" % d.root) for d in hse.loc_indices}
# Only retain the halos required by the Diag scheme
# Note: `sorted` is only for deterministic code generation
halos = sorted(i for i in hse.halos if isinstance(i.dim, tuple))
body = []
for dims, tosides in halos:
mapper = OrderedDict(zip(dims, [i.flip() for i in tosides]))
fromrank = FieldFromPointer(
''.join(i.name[0] for i in mapper.values()), nb)
ofss = [
fixed.get(d,
f._C_get_field(HALO, d, mapper.get(d)).offset)
for d in f.dimensions
]
msgi = Byref(IndexedPointer(msg, len(body)))
body.append(Call(wait.name, [f] + ofss + [fromrank, msgi]))
iet = List(body=body)
parameters = [f] + list(fixed.values()) + [nb, msg]
return Callable('halowait%d' % key, iet, 'void', parameters,
('static', ))
def _call_sendrecv(self, name, *args, msg=None, haloid=None):
# Drop `sizes` as this HaloExchangeBuilder conveys them through `msg`
# Drop `ofss` as this HaloExchangeBuilder only needs them in `wait()`,
# to collect and scatter the result of an MPI_Irecv
f, _, ofsg, _, fromrank, torank, comm = args
msg = Byref(IndexedPointer(msg, haloid))
return Call(name, [f] + ofsg + [fromrank, torank, comm, msg])
def _make_halowait(self, f, hse, key, msg=None):
cast = cast_mapper[(f.dtype, '*')]
fixed = {d: Symbol(name="o%s" % d.root) for d in hse.loc_indices}
dim = Dimension(name='i')
msgi = IndexedPointer(msg, dim)
bufs = FieldFromComposite(msg._C_field_bufs, msgi)
fromrank = FieldFromComposite(msg._C_field_from, msgi)
sizes = [FieldFromComposite('%s[%d]' % (msg._C_field_sizes, i), msgi)
for i in range(len(f._dist_dimensions))]
ofss = [FieldFromComposite('%s[%d]' % (msg._C_field_ofss, i), msgi)
for i in range(len(f._dist_dimensions))]
ofss = [fixed.get(d) or ofss.pop(0) for d in f.dimensions]
# The `scatter` must be guarded as we must not alter the halo values along
# the domain boundary, where the sender is actually MPI.PROC_NULL
scatter = Call('scatter%s' % key, [cast(bufs)] + sizes + [f] + ofss)
scatter = Conditional(CondNe(fromrank, Macro('MPI_PROC_NULL')), scatter)
rrecv = Byref(FieldFromComposite(msg._C_field_rrecv, msgi))
waitrecv = Call('MPI_Wait', [rrecv, Macro('MPI_STATUS_IGNORE')])
rsend = Byref(FieldFromComposite(msg._C_field_rsend, msgi))
waitsend = Call('MPI_Wait', [rsend, Macro('MPI_STATUS_IGNORE')])
# The -1 below is because an Iteration, by default, generates <=
ncomms = Symbol(name='ncomms')
iet = Iteration([waitsend, waitrecv, scatter], dim, ncomms - 1)
parameters = ([f] + list(fixed.values()) + [msg, ncomms])
return Callable('halowait%d' % key, iet, 'void', parameters, ('static',))
def _make_haloupdate(self, f, hse, key, msg=None):
comm = f.grid.distributor._obj_comm
fixed = {d: Symbol(name="o%s" % d.root) for d in hse.loc_indices}
dim = Dimension(name='i')
msgi = IndexedPointer(msg, dim)
bufg = FieldFromComposite(msg._C_field_bufg, msgi)
bufs = FieldFromComposite(msg._C_field_bufs, msgi)
fromrank = FieldFromComposite(msg._C_field_from, msgi)
torank = FieldFromComposite(msg._C_field_to, msgi)
sizes = [
FieldFromComposite('%s[%d]' % (msg._C_field_sizes, i), msgi)
for i in range(len(f._dist_dimensions))
]
ofsg = [
FieldFromComposite('%s[%d]' % (msg._C_field_ofsg, i), msgi)
for i in range(len(f._dist_dimensions))
]
ofsg = [fixed.get(d) or ofsg.pop(0) for d in f.dimensions]
# The `gather` is unnecessary if sending to MPI.PROC_NULL
gather = Call('gather_%s' % key, [bufg] + sizes + [f] + ofsg)
gather = Conditional(CondNe(torank, Macro('MPI_PROC_NULL')), gather)
# Make Irecv/Isend
count = reduce(mul, sizes, 1)
rrecv = Byref(FieldFromComposite(msg._C_field_rrecv, msgi))
rsend = Byref(FieldFromComposite(msg._C_field_rsend, msgi))
recv = Call('MPI_Irecv', [
bufs, count,
Macro(dtype_to_mpitype(f.dtype)), fromrank,
Integer(13), comm, rrecv
])
send = Call('MPI_Isend', [
bufg, count,
Macro(dtype_to_mpitype(f.dtype)), torank,
Integer(13), comm, rsend
])
# The -1 below is because an Iteration, by default, generates <=
ncomms = Symbol(name='ncomms')
iet = Iteration([recv, gather, send], dim, ncomms - 1)
parameters = ([f, comm, msg, ncomms]) + list(fixed.values())
return Callable('haloupdate%d' % key, iet, 'void', parameters,
('static', ))
def _make_remainder(self, hs, key, callcompute, region):
assert callcompute.is_Call
dim = Dimension(name='i')
region_i = IndexedPointer(region, dim)
dynamic_args_mapper = {}
for i in hs.arguments:
if i.is_Dimension:
dynamic_args_mapper[i] = (FieldFromComposite(i.min_name, region_i),
FieldFromComposite(i.max_name, region_i))
else:
dynamic_args_mapper[i] = (FieldFromComposite(i.name, region_i),)
iet = callcompute._rebuild(dynamic_args_mapper=dynamic_args_mapper)
# The -1 below is because an Iteration, by default, generates <=
iet = Iteration(iet, dim, region.nregions - 1)
return make_efunc('remainder%d' % key, iet)
def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage):
"""
Allocate the following objects in the high bandwidth memory:
* The pointer array `obj`;
* The pointee Array `obj.array`
If the pointer array is defined over `sregistry.threadid`, that is a thread
Dimension, then each `obj.array` slice is allocated and freed individually
by the owner thread.
"""
# The pointer array
decl = Definition(obj)
memptr = VOID(Byref(obj._C_symbol), '**')
alignment = obj._data_alignment
size = SizeOf(Keyword('%s*' % obj._C_typedata)) * obj.dim.symbolic_size
alloc0 = self.lang['host-alloc'](memptr, alignment, size)
free0 = self.lang['host-free'](obj._C_symbol)
# The pointee Array
pobj = IndexedPointer(obj._C_symbol, obj.dim)
memptr = VOID(Byref(pobj), '**')
size = SizeOf(obj._C_typedata) * prod(obj.array.symbolic_shape)
alloc1 = self.lang['host-alloc'](memptr, alignment, size)
free1 = self.lang['host-free'](pobj)
# Dump
if obj.dim is self.sregistry.threadid:
storage.update(obj,
site,
allocs=(decl, alloc0),
frees=free0,
pallocs=(obj.dim, alloc1),
pfrees=(obj.dim, free1))
else:
storage.update(obj,
site,
allocs=(decl, alloc0, alloc1),
frees=(free0, free1))
def _make_poke(self, hs, key, msgs):
flag = Symbol(name='flag')
initflag = LocalExpression(DummyEq(flag, 0))
body = [initflag]
for msg in msgs:
dim = Dimension(name='i')
msgi = IndexedPointer(msg, dim)
rrecv = Byref(FieldFromComposite(msg._C_field_rrecv, msgi))
rsend = Byref(FieldFromComposite(msg._C_field_rsend, msgi))
testrecv = Call(
'MPI_Test',
[rrecv, Byref(flag),
Macro('MPI_STATUS_IGNORE')])
testsend = Call(
'MPI_Test',
[rsend, Byref(flag),
Macro('MPI_STATUS_IGNORE')])
body.append(Iteration([testsend, testrecv], dim, msg.npeers - 1))
return make_efunc('pokempi%d' % key, body)