def freesurface(model, eq):
"""
Generate the stencil that mirrors the field as a free surface modeling for
the acoustic wave equation.
Parameters
----------
model : Model
Physical model.
eq : Eq
Time-stepping stencil (time update) to mirror at the freesurface.
"""
lhs, rhs = eq.evaluate.args
# Get vertical dimension and corresponding subdimension
zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
z = zfs.parent
# Functions present in the stencil
funcs = retrieve_functions(rhs)
mapper = {}
# Antisymmetric mirror at negative indices
# TODO: Make a proper "mirror_indices" tool function
for f in funcs:
zind = f.indices[-1]
if (zind - z).as_coeff_Mul()[0] < 0:
s = sign(zind.subs({z: zfs, z.spacing: 1}))
mapper.update({f: s * f.subs({zind: INT(abs(zind))})})
return Eq(lhs, rhs.subs(mapper), subdomain=model.grid.subdomains['fsdomain'])
def inject(self, field, expr, offset=0):
"""
Generate equations injecting an arbitrary expression into a field.
Parameters
----------
field : Function
Input field into which the injection is performed.
expr : expr-like
Injected expression.
offset : int, optional
Additional offset from the boundary.
"""
expr = indexify(expr)
field = indexify(field)
p, _ = self.gridpoints.indices
dim_subs = []
coeffs = []
for i, d in enumerate(self.grid.dimensions):
rd = DefaultDimension(name="r%s" % d.name, default_value=self.r)
dim_subs.append((d, INT(rd + self.gridpoints[p, i])))
coeffs.append(self.interpolation_coeffs[p, i, rd])
rhs = prod(coeffs) * expr
field = field.subs(dim_subs)
return [Eq(field, field + rhs.subs(dim_subs))]
def _coordinate_indices(self):
"""Symbol for each grid index according to the coordinates."""
return tuple([
INT(FLOOR((c - o) / i.spacing))
for c, o, i in zip(self._coordinate_symbols, self.grid.origin,
self.grid.dimensions[:self.grid.dim])
])
def interpolate(self, expr, offset=0, increment=False, self_subs={}):
"""
Generate equations interpolating an arbitrary expression into ``self``.
Parameters
----------
expr : expr-like
Input expression to interpolate.
offset : int, optional
Additional offset from the boundary.
increment: bool, optional
If True, generate increments (Inc) rather than assignments (Eq).
"""
expr = indexify(expr)
p, _, _ = self.interpolation_coeffs.indices
dim_subs = []
coeffs = []
for i, d in enumerate(self.grid.dimensions):
rd = DefaultDimension(name="r%s" % d.name, default_value=self.r)
dim_subs.append((d, INT(rd + self.gridpoints[p, i])))
coeffs.append(self.interpolation_coeffs[p, i, rd])
# Apply optional time symbol substitutions to lhs of assignment
lhs = self.subs(self_subs)
rhs = prod(coeffs) * expr.subs(dim_subs)
return [Eq(lhs, lhs + rhs)]
def _coordinate_indices(self):
"""Symbol for each grid index according to the coordinates."""
indices = self.grid.dimensions
return tuple([
INT(sympy.Function('floor')((c - o) / i.spacing))
for c, o, i in zip(self._coordinate_symbols, self.grid.origin,
indices[:self.grid.dim])
])
def make_var_gets(expr):
mapper = {}
indexeds = retrieve_indexed(expr)
data_carriers = [i for i in indexeds if i.base.function.from_YASK]
for i in data_carriers:
args = [
ListInitializer([INT(make_var_gets(j)) for j in i.indices])
]
mapper[i] = make_sharedptr_funcall(
namespace['code-var-get'], args,
yk_var_objs[i.base.function.name])
return expr.xreplace(mapper)
def _make_partree(self, candidates, nthreads=None):
"""Parallelize the `candidates` Iterations attaching suitable OpenMP pragmas."""
assert candidates
root = candidates[0]
# Get the collapsable Iterations
collapsable = self._find_collapsable(root, candidates)
ncollapse = 1 + len(collapsable)
# Prepare to build a ParallelTree
if all(i.is_Affine for i in candidates):
bundles = FindNodes(ExpressionBundle).visit(root)
sops = sum(i.ops for i in bundles)
if sops >= self.dynamic_work:
schedule = 'dynamic'
else:
schedule = 'static'
if nthreads is None:
# pragma omp for ... schedule(..., 1)
nthreads = self.nthreads
body = OpenMPIteration(schedule=schedule,
ncollapse=ncollapse,
**root.args)
else:
# pragma omp parallel for ... schedule(..., 1)
body = OpenMPIteration(schedule=schedule,
parallel=True,
ncollapse=ncollapse,
nthreads=nthreads,
**root.args)
prefix = []
else:
# pragma omp for ... schedule(..., expr)
assert nthreads is None
nthreads = self.nthreads_nonaffine
chunk_size = Symbol(name='chunk_size')
body = OpenMPIteration(ncollapse=ncollapse,
chunk_size=chunk_size,
**root.args)
niters = prod([root.symbolic_size] +
[j.symbolic_size for j in collapsable])
value = INT(Max(niters / (nthreads * self.chunk_nonaffine), 1))
prefix = [Expression(DummyEq(chunk_size, value, dtype=np.int32))]
# Create a ParallelTree
partree = ParallelTree(prefix, body, nthreads=nthreads)
collapsed = [partree] + collapsable
return root, partree, collapsed
def callback():
_expr = indexify(expr)
_field = indexify(field)
p, _ = self.obj.gridpoints.indices
dim_subs = []
coeffs = []
for i, d in enumerate(self.obj.grid.dimensions):
rd = DefaultDimension(name="r%s" % d.name, default_value=self.r)
dim_subs.append((d, INT(rd + self.obj.gridpoints[p, i])))
coeffs.append(self.obj.interpolation_coeffs[p, i, rd])
rhs = prod(coeffs) * _expr
_field = _field.subs(dim_subs)
return [Eq(_field, _field + rhs.subs(dim_subs))]
def make_var_accesses(node, yk_var_objs):
"""
Construct a new Iteration/Expression based on ``node``, in which all
:class:`types.Indexed` accesses have been converted into YASK var
accesses.
"""
def make_var_gets(expr):
mapper = {}
indexeds = retrieve_indexed(expr)
data_carriers = [i for i in indexeds if i.base.function.from_YASK]
for i in data_carriers:
args = [
ListInitializer([INT(make_var_gets(j)) for j in i.indices])
]
mapper[i] = make_sharedptr_funcall(
namespace['code-var-get'], args,
yk_var_objs[i.base.function.name])
return expr.xreplace(mapper)
mapper = {}
for i, e in enumerate(FindNodes(Expression).visit(node)):
if e.is_ForeignExpression:
continue
lhs, rhs = e.expr.args
# RHS translation
rhs = make_var_gets(rhs)
# LHS translation
if e.write.from_YASK:
args = [rhs]
args += [
ListInitializer([INT(make_var_gets(i)) for i in lhs.indices])
]
call = namespace['code-var-add' if e.
is_Increment else 'code-var-put']
handle = make_sharedptr_funcall(call, args,
yk_var_objs[e.write.name])
processed = ForeignExpression(handle,
e.dtype,
is_Increment=e.is_Increment)
else:
# Writing to a scalar temporary
processed = e._rebuild(expr=e.expr.func(lhs, rhs))
mapper.update({e: processed})
return Transformer(mapper).visit(node)
def callback():
_expr = indexify(expr)
p, _, _ = self.obj.interpolation_coeffs.indices
dim_subs = []
coeffs = []
for i, d in enumerate(self.obj.grid.dimensions):
rd = DefaultDimension(name="r%s" % d.name, default_value=self.r)
dim_subs.append((d, INT(rd + self.obj.gridpoints[p, i])))
coeffs.append(self.obj.interpolation_coeffs[p, i, rd])
# Apply optional time symbol substitutions to lhs of assignment
lhs = self.obj.subs(self_subs)
rhs = prod(coeffs) * _expr.subs(dim_subs)
return [Eq(lhs, lhs + rhs)]
def _make_partree(self, candidates, nthreads=None):
"""Parallelize the `candidates` Iterations attaching suitable OpenMP pragmas."""
assert candidates
root = candidates[0]
# Get the collapsable Iterations
collapsable = self._find_collapsable(root, candidates)
ncollapse = 1 + len(collapsable)
# 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
def freesurface(model, eq):
"""
Generate the stencil that mirrors the field as a free surface modeling for
the acoustic wave equation
Parameters
----------
model: Model
Physical model
eq: Eq or List of Eq
Equation to apply mirror to
"""
fs_eq = []
for eq_i in eq:
for p in eq_i._flatten:
lhs, rhs = p.evaluate.args
# Add modulo replacements to to rhs
zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
z = zfs.parent
funcs = retrieve_functions(rhs.evaluate)
mapper = {}
for f in funcs:
zind = f.indices[-1]
if (zind - z).as_coeff_Mul()[0] < 0:
s = sign((zind - z.symbolic_min).subs({
z: zfs,
z.spacing: 1
}))
mapper.update({f: s * f.subs({zind: INT(abs(zind))})})
fs_eq.append(
Eq(lhs,
sign(lhs.indices[-1] - z.symbolic_min) * rhs.subs(mapper),
subdomain=model.grid.subdomains['fsdomain']))
return fs_eq
def _make_partree(self, candidates, nthreads=None):
"""Parallelize `root` attaching a suitable OpenMP pragma."""
assert candidates
root = candidates[0]
# Get the collapsable Iterations
collapsable = []
if ncores() >= Ompizer.COLLAPSE_NCORES and IsPerfectIteration().visit(
root):
for n, i in enumerate(candidates[1:], 1):
# The OpenMP specification forbids collapsed loops to use iteration
# variables in initializer expressions. E.g., the following is forbidden:
#
# #pragma omp ... collapse(2)
# for (i = ... )
# for (j = i ...)
# ...
#
# Here, we make sure this won't happen
if any(j.dim in i.symbolic_min.free_symbols
for j in candidates[:n]):
break
# Also, we do not want to collapse vectorizable Iterations
if i.is_Vectorizable:
break
# Would there be enough work per parallel iteration?
try:
work = prod(
[int(j.dim.symbolic_size) for j in candidates[n + 1:]])
if work < Ompizer.COLLAPSE_WORK:
break
except TypeError:
pass
collapsable.append(i)
ncollapse = 1 + len(collapsable)
# 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
