def yaskit(trees, yc_soln):
"""
Populate a YASK compiler solution with the :class:`Expression`s found in an IET.
The necessary YASK grids are instantiated.
:param trees: A sequence of offloadable :class:`IterationTree`s, in which the
Expressions are searched.
:param yc_soln: The YASK compiler solution to be populated.
"""
# Track all created YASK grids
mapper = {}
# It's up to Devito to organize the equations into a flow graph
yc_soln.set_dependency_checker_enabled(False)
processed = []
for tree in trees:
# All expressions within `tree`
expressions = [i.expr for i in FindNodes(Expression).visit(tree.inner)]
# Attach conditional expression for sub-domains
conditions = [(i, []) for i in expressions]
for i in tree:
if not i.dim.is_Sub:
continue
# Can we express both Iteration extremes as
# `FIRST(i.dim) + integer` OR `LAST(i.dim) + integer` ?
# If not, one of the following lines will throw a TypeError exception
lower_ofs, lower_sym = i.dim.offset_left()
upper_ofs, upper_sym = i.dim.offset_right()
if i.is_Parallel:
# At this point, no issues are expected -- we should just be able to
# build the YASK conditions under which to execute this parallel Iteration
ydim = nfac.new_domain_index(i.dim.parent.name)
# Handle lower extreme
if lower_sym == i.dim.parent.symbolic_start:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
expr = nfac.new_add_node(node,
nfac.new_const_number_node(lower_ofs))
for _, v in conditions:
v.append(nfac.new_not_less_than_node(ydim, expr))
# Handle upper extreme
if upper_sym == i.dim.parent.symbolic_start:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
expr = nfac.new_add_node(node,
nfac.new_const_number_node(upper_ofs))
for _, v in conditions:
v.append(nfac.new_not_greater_than_node(ydim, expr))
elif i.is_Sequential:
# For sequential Iterations, the extent *must* be statically known,
# otherwise we don't know how to handle this
try:
int(i.extent())
except TypeError:
raise NotImplementedError(
"Found sequential Iteration with "
"statically unknown extent")
assert lower_sym == upper_sym # A corollary of getting up to this point
n = lower_sym
ydim = nfac.new_domain_index(i.dim.parent.name)
if n == i.dim.parent.symbolic_start:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
if i.direction is Backward:
_range = range(upper_ofs, lower_ofs - 1, -1)
else:
_range = range(lower_ofs, upper_ofs + 1)
unwound = []
for e, v in conditions:
for r in _range:
expr = nfac.new_add_node(node,
nfac.new_const_number_node(r))
unwound.append(
(e, v + [nfac.new_equals_node(ydim, expr)]))
conditions = unwound
# Build the YASK equations as well as all necessary grids
for k, v in conditions:
yask_expr = make_yask_ast(k, yc_soln, mapper)
if yask_expr is not None:
processed.append(yask_expr)
# Is there a sub-domain to attach ?
if v:
condition = v.pop(0)
for i in v:
condition = nfac.new_and_node(condition, i)
yask_expr.set_cond(condition)
# Add flow dependences to the offloaded equations
# TODO: This can be improved by spotting supergroups ?
for to, frm in zip(processed, processed[1:]):
yc_soln.add_flow_dependency(frm, to)
return mapper
def yaskit(trees, yc_soln):
"""
Populate a YASK compiler solution with the Expressions found in an IET.
The necessary YASK vars are instantiated.
Parameters
----------
trees : list of IterationTree
One or more Iteration trees within which the convertible Expressions are searched.
yc_soln
The YASK compiler solution to be populated.
"""
# It's up to Devito to organize the equations into a flow graph
yc_soln.set_dependency_checker_enabled(False)
mapper = {}
processed = []
for tree in trees:
# All expressions within `tree`
expressions = [i.expr for i in FindNodes(Expression).visit(tree.inner)]
# Attach conditional expression for sub-domains
conditions = [(i, []) for i in expressions]
for i in tree:
if not i.dim.is_Sub:
continue
# Can we express both Iteration extremes as
# `FIRST(i.dim) + integer` OR `LAST(i.dim) + integer` ?
# If not, one of the following lines will throw a TypeError exception
lvalue, lextreme, _ = i.dim._offset_left
rvalue, rextreme, _ = i.dim._offset_right
if i.is_Parallel:
# At this point, no issues are expected -- we should just be able to
# build the YASK conditions under which to execute this parallel Iteration
ydim = nfac.new_domain_index(i.dim.parent.name)
# Handle lower extreme
if lextreme == i.dim.parent.symbolic_min:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
expr = nfac.new_add_node(node,
nfac.new_const_number_node(lvalue))
for _, v in conditions:
v.append(nfac.new_not_less_than_node(ydim, expr))
# Handle upper extreme
if rextreme == i.dim.parent.symbolic_min:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
expr = nfac.new_add_node(node,
nfac.new_const_number_node(rvalue))
for _, v in conditions:
v.append(nfac.new_not_greater_than_node(ydim, expr))
elif i.is_Sequential:
# For sequential Iterations, the extent *must* be statically known,
# otherwise we don't know how to handle this
try:
int(i.dim._thickness_map.get(i.size()))
except TypeError:
raise NotImplementedError(
"Found sequential Iteration with "
"statically unknown extent")
assert lextreme == rextreme # A corollary of getting up to this point
n = lextreme
ydim = nfac.new_domain_index(i.dim.parent.name)
if n == i.dim.parent.symbolic_min:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
if i.direction is Backward:
_range = range(rvalue, lvalue - 1, -1)
else:
_range = range(lvalue, rvalue + 1)
unwound = []
for e, v in conditions:
for r in _range:
expr = nfac.new_add_node(node,
nfac.new_const_number_node(r))
unwound.append(
(e, v + [nfac.new_equals_node(ydim, expr)]))
conditions = unwound
# Build the YASK equations as well as all necessary vars
for k, v in conditions:
yask_expr = make_yask_ast(k, yc_soln, mapper)
if yask_expr is not None:
processed.append(yask_expr)
# Is there a sub-domain to attach ?
if v:
condition = v.pop(0)
for i in v:
condition = nfac.new_and_node(condition, i)
yask_expr.set_cond(condition)
# Add flow dependences to the offloaded equations
# TODO: This can be improved by spotting supergroups ?
for to, frm in zip(processed, processed[1:]):
yc_soln.add_flow_dependency(frm, to)
# Have we built any local vars (i.e., DSE-produced tensor temporaries)?
# If so, eventually these will be mapped to YASK scratch vars
local_vars = [i for i in mapper if i.is_Array]
return local_vars
def yaskit(trees, yc_soln):
"""
Populate a YASK compiler solution with the :class:`Expression`s found in an IET.
The necessary YASK grids are instantiated.
Parameters
----------
trees : list of IterationTree
One or more Iteration trees within which the convertible Expressions are searched.
yc_soln
The YASK compiler solution to be populated.
"""
# It's up to Devito to organize the equations into a flow graph
yc_soln.set_dependency_checker_enabled(False)
mapper = {}
processed = []
for tree in trees:
# All expressions within `tree`
expressions = [i.expr for i in FindNodes(Expression).visit(tree.inner)]
# Attach conditional expression for sub-domains
conditions = [(i, []) for i in expressions]
for i in tree:
if not i.dim.is_Sub:
continue
# Can we express both Iteration extremes as
# `FIRST(i.dim) + integer` OR `LAST(i.dim) + integer` ?
# If not, one of the following lines will throw a TypeError exception
lower_ofs, lower_sym = i.dim._offset_left()
upper_ofs, upper_sym = i.dim._offset_right()
if i.is_Parallel:
# At this point, no issues are expected -- we should just be able to
# build the YASK conditions under which to execute this parallel Iteration
ydim = nfac.new_domain_index(i.dim.parent.name)
# Handle lower extreme
if lower_sym == i.dim.parent.symbolic_min:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
expr = nfac.new_add_node(node, nfac.new_const_number_node(lower_ofs))
for _, v in conditions:
v.append(nfac.new_not_less_than_node(ydim, expr))
# Handle upper extreme
if upper_sym == i.dim.parent.symbolic_min:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
expr = nfac.new_add_node(node, nfac.new_const_number_node(upper_ofs))
for _, v in conditions:
v.append(nfac.new_not_greater_than_node(ydim, expr))
elif i.is_Sequential:
# For sequential Iterations, the extent *must* be statically known,
# otherwise we don't know how to handle this
try:
int(i.dim._thickness_map.get(i.size()))
except TypeError:
raise NotImplementedError("Found sequential Iteration with "
"statically unknown extent")
assert lower_sym == upper_sym # A corollary of getting up to this point
n = lower_sym
ydim = nfac.new_domain_index(i.dim.parent.name)
if n == i.dim.parent.symbolic_min:
node = nfac.new_first_domain_index(ydim)
else:
node = nfac.new_last_domain_index(ydim)
if i.direction is Backward:
_range = range(upper_ofs, lower_ofs - 1, -1)
else:
_range = range(lower_ofs, upper_ofs + 1)
unwound = []
for e, v in conditions:
for r in _range:
expr = nfac.new_add_node(node, nfac.new_const_number_node(r))
unwound.append((e, v + [nfac.new_equals_node(ydim, expr)]))
conditions = unwound
# Build the YASK equations as well as all necessary grids
for k, v in conditions:
yask_expr = make_yask_ast(k, yc_soln, mapper)
if yask_expr is not None:
processed.append(yask_expr)
# Is there a sub-domain to attach ?
if v:
condition = v.pop(0)
for i in v:
condition = nfac.new_and_node(condition, i)
yask_expr.set_cond(condition)
# Add flow dependences to the offloaded equations
# TODO: This can be improved by spotting supergroups ?
for to, frm in zip(processed, processed[1:]):
yc_soln.add_flow_dependency(frm, to)
# Have we built any local grids (i.e., DSE-produced tensor temporaries)?
# If so, eventually these will be mapped to YASK scratch grids
local_grids = [i for i in mapper if i.is_Array]
return local_grids