Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
@pytest.mark.parametrize('mode,expected', [
('free-symbols', '["f", "x"]'),
('symbolics', '["f"]')
])
def test_find_symbols_nested(self, mode, expected):
grid = Grid(shape=(4, 4, 4))
call = Call('foo', [
Call('bar', [
Symbol(name='x'),
Call('baz', [Function(name='f', grid=grid)])
])
])
found = FindSymbols(mode).visit(call)
assert [f.name for f in found] == eval(expected)
def test_to_ops_stencil(self, _accesses):
param = Symbol('foo')
accesses = eval(_accesses)
stencil_name = 's2d_foo_%spt' % len(accesses)
stencil, result = to_ops_stencil(param, accesses)
assert stencil.name == stencil_name.upper()
assert result[0].expr.lhs.name == stencil_name
assert result[0].expr.rhs.params == tuple(itertools.chain(*accesses))
assert result[1].expr.lhs == stencil
assert type(result[1].expr.rhs) == namespace['ops_decl_stencil']
assert result[1].expr.rhs.args == (
2,
len(accesses),
def _specialize_iet(self, iet, **kwargs):
in1 = Symbol(name='in1')
in2 = Symbol(name='in2')
out1 = Symbol(name='out1')
out2 = Symbol(name='out2')
# extract expression from iet
expr_in = FindNodes(Expression).visit(iet)[0]
equation = expr_in.expr
# extract symbols from equation
symb = retrieve_indexed(equation)
symb_ordered = [symb[1], symb[2], symb[0]] # TODO: func to order symb
symb_ordered_ext = [in2, in1] + symb_ordered + [out1, out2]
# draw successive symbol shifts
# first pass
mapping = []
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')])
seen = {}
for i, e in enumerate(exprs):
seen.setdefault(e.lhs, []).append(i)
# Optimization: don't waste time reconstructing stuff if already in SSA form
if all(len(i) == 1 for i in seen.values()):
return exprs
# SSA conversion
c = 0
mapper = {}
processed = []
for i, e in enumerate(exprs):
where = seen[e.lhs]
rhs = e.rhs.xreplace(mapper)
if len(where) > 1:
needssa = e.is_Scalar or where[-1] != i
lhs = Symbol(name='ssa%d' % c, dtype=e.dtype) if needssa else e.lhs
if e.is_Increment:
# Turn AugmentedAssignment into Assignment
processed.append(e.func(lhs, mapper[e.lhs] + rhs, is_Increment=False))
else:
processed.append(e.func(lhs, rhs))
mapper[e.lhs] = lhs
c += 1
else:
processed.append(e.func(e.lhs, rhs))
return processed
dtype=np.int32,
)
ops_stencil = OpsStencil(stencil_name.upper())
return ops_stencil, [
Expression(ClusterizedEq(Eq(
stencil_array,
ListInitializer(list(itertools.chain(*accesses)))
))),
Expression(ClusterizedEq(Eq(
ops_stencil,
namespace['ops_decl_stencil'](
dims,
pts,
Symbol(stencil_array.name),
Literal('"%s"' % stencil_name.upper())
)
def _make_copy(self, f, hse, key, swap=False):
buf_dims = []
buf_indices = []
for d in f.dimensions:
if d not in hse.loc_indices:
buf_dims.append(Dimension(name='buf_%s' % d.root))
buf_indices.append(d.root)
buf = Array(name='buf', dimensions=buf_dims, dtype=f.dtype, padding=0)
f_offsets = []
f_indices = []
for d in f.dimensions:
offset = Symbol(name='o%s' % d.root)
f_offsets.append(offset)
f_indices.append(offset + (d.root if d not in hse.loc_indices else 0))
if swap is False:
eq = DummyEq(buf[buf_indices], f[f_indices])
name = 'gather_%s' % key
else:
eq = DummyEq(f[f_indices], buf[buf_indices])
name = 'scatter_%s' % key
iet = Expression(eq)
for i, d in reversed(list(zip(buf_indices, buf_dims))):
# The -1 below is because an Iteration, by default, generates <=
iet = Iteration(iet, i, d.symbolic_size - 1, properties=(PARALLEL, AFFINE))
parameters = [buf] + list(buf.shape) + [f] + f_offsets
def _make_sendrecv(self, f, hse, key, **kwargs):
comm = f.grid.distributor._obj_comm
buf_dims = [Dimension(name='buf_%s' % d.root) for d in f.dimensions
if d not in hse.loc_indices]
bufg = Array(name='bufg', dimensions=buf_dims, dtype=f.dtype,
padding=0, scope='heap')
bufs = Array(name='bufs', dimensions=buf_dims, dtype=f.dtype,
padding=0, scope='heap')
ofsg = [Symbol(name='og%s' % d.root) for d in f.dimensions]
ofss = [Symbol(name='os%s' % d.root) for d in f.dimensions]
fromrank = Symbol(name='fromrank')
torank = Symbol(name='torank')
gather = Call('gather_%s' % key, [bufg] + list(bufg.shape) + [f] + ofsg)
scatter = Call('scatter_%s' % key, [bufs] + list(bufs.shape) + [f] + ofss)
# The `gather` is unnecessary if sending to MPI.PROC_NULL
gather = Conditional(CondNe(torank, Macro('MPI_PROC_NULL')), gather)
# 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 = Conditional(CondNe(fromrank, Macro('MPI_PROC_NULL')), scatter)
count = reduce(mul, bufs.shape, 1)
rrecv = MPIRequestObject(name='rrecv')
rsend = MPIRequestObject(name='rsend')
recv = Call('MPI_Irecv', [bufs, count, Macro(dtype_to_mpitype(f.dtype)),
fromrank, Integer(13), comm, rrecv])
def _make_sendrecv(self, f, hse, key, msg=None):
comm = f.grid.distributor._obj_comm
bufg = FieldFromPointer(msg._C_field_bufg, msg)
bufs = FieldFromPointer(msg._C_field_bufs, msg)
ofsg = [Symbol(name='og%s' % d.root) for d in f.dimensions]
fromrank = Symbol(name='fromrank')
torank = Symbol(name='torank')
sizes = [FieldFromPointer('%s[%d]' % (msg._C_field_sizes, i), msg)
for i in range(len(f._dist_dimensions))]
gather = Call('gather_%s' % key, [bufg] + sizes + [f] + ofsg)
# The `gather` is unnecessary if sending to MPI.PROC_NULL
gather = Conditional(CondNe(torank, Macro('MPI_PROC_NULL')), gather)
count = reduce(mul, sizes, 1)
rrecv = Byref(FieldFromPointer(msg._C_field_rrecv, msg))
rsend = Byref(FieldFromPointer(msg._C_field_rsend, msg))
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])
# Prepare to build a ParallelTree
prefix = []
if all(i.is_Affine for i in candidates):
if nthreads is None:
# pragma omp for ... schedule(..., 1)
nthreads = self.nthreads
omp_pragma = self.lang['for'](ncollapse, 1)
else:
# pragma omp parallel for ... schedule(..., 1)
omp_pragma = self.lang['par-for'](ncollapse, 1, nthreads)
else:
# pragma omp for ... schedule(..., expr)
assert nthreads is None
nthreads = self.nthreads_nonaffine
chunk_size = Symbol(name='chunk_size')
omp_pragma = self.lang['for'](ncollapse, chunk_size)
niters = prod([root.symbolic_size] + [j.symbolic_size for j in collapsable])
value = INT(Max(niters / (nthreads*self.CHUNKSIZE_NONAFFINE), 1))
prefix.append(Expression(DummyEq(chunk_size, value, dtype=np.int32)))
# Create a ParallelTree
body = root._rebuild(pragmas=root.pragmas + (omp_pragma,),
properties=root.properties + (COLLAPSED(ncollapse),))
partree = ParallelTree(prefix, body, nthreads=nthreads)
collapsed = [partree] + collapsable
return root, partree, collapsed