def make_grid_accesses(node, yk_grid_objs):
"""
Construct a new Iteration/Expression based on ``node``, in which all
:class:`types.Indexed` accesses have been converted into YASK grid
accesses.
"""
def make_grid_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_grid_gets(j)) for j in i.indices])
]
mapper[i] = make_sharedptr_funcall(
namespace['code-grid-get'], args,
yk_grid_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_grid_gets(rhs)
# LHS translation
if e.write.from_YASK:
args = [rhs]
args += [
ListInitializer([INT(make_grid_gets(i)) for i in lhs.indices])
]
call = namespace['code-grid-add' if e.
is_Increment else 'code-grid-put']
handle = make_sharedptr_funcall(call, args,
yk_grid_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 _specialize_iet(self, iet, **kwargs):
"""
Transform the Iteration/Expression tree to offload the computation of
one or more loop nests onto YASK. This involves calling the YASK compiler
to generate YASK code. Such YASK code is then called from within the
transformed Iteration/Expression tree.
"""
mapper = {}
self.yk_solns = OrderedDict()
for n, (section, trees) in enumerate(find_affine_trees(iet).items()):
dimensions = tuple(
filter_ordered(i.dim.root for i in flatten(trees)))
context = contexts.fetch(dimensions, self._dtype)
# A unique name for the 'real' compiler and kernel solutions
name = namespace['jit-soln'](Signer._digest(
configuration, *[i.root for i in trees]))
# Create a YASK compiler solution for this Operator
yc_soln = context.make_yc_solution(name)
try:
# Generate YASK grids and populate `yc_soln` with equations
local_grids = yaskit(trees, yc_soln)
# Build the new IET nodes
yk_soln_obj = YaskSolnObject(namespace['code-soln-name'](n))
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'], yk_soln_obj)
funcall = ForeignExpression(funcall, self._dtype)
mapper[trees[0].root] = funcall
mapper.update({i.root: mapper.get(i.root)
for i in trees}) # Drop trees
# Mark `funcall` as an external function call
self._func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
yk_soln = context.make_yk_solution(name, yc_soln, local_grids)
self.yk_solns[(dimensions, yk_soln_obj)] = yk_soln
# Print some useful information about the newly constructed solution
log("Solution '%s' contains %d grid(s) and %d equation(s)." %
(yc_soln.get_name(), yc_soln.get_num_grids(),
yc_soln.get_num_equations()))
except NotImplementedError as e:
log("Unable to offload a candidate tree. Reason: [%s]" %
str(e))
iet = Transformer(mapper).visit(iet)
if not self.yk_solns:
log("No offloadable trees found")
# Some Iteration/Expression trees are not offloaded to YASK and may
# require further processing to be executed in YASK, due to the differences
# in storage layout employed by Devito and YASK
yk_grid_objs = {
i.name: YaskGridObject(i.name)
for i in self._input if i.from_YASK
}
yk_grid_objs.update({i: YaskGridObject(i) for i in self._local_grids})
iet = make_grid_accesses(iet, yk_grid_objs)
# Finally optimize all non-yaskized loops
iet = super(OperatorYASK, self)._specialize_iet(iet, **kwargs)
return iet