def run(expr):
if expr.is_Integer:
return nfac.new_const_number_node(int(expr))
elif expr.is_Float:
return nfac.new_const_number_node(float(expr))
elif expr.is_Symbol:
function = expr.base.function
if function.is_Constant:
if function not in self.mapper:
self.mapper[function] = self.yc_soln.new_grid(
function.name, [])
return self.mapper[function].new_relative_grid_point([])
else:
# A DSE-generated temporary, which must have already been
# encountered as a LHS of a previous expression
assert function in self.mapper
return self.mapper[function]
elif isinstance(expr, Indexed):
function = expr.base.function
if function not in self.mapper:
if function.is_TimeFunction:
dimensions = [
nfac.new_step_index(function.indices[0].name)
]
dimensions += [
nfac.new_domain_index(i.name)
for i in function.indices[1:]
]
else:
dimensions = [
nfac.new_domain_index(i.name)
for i in function.indices
]
self.mapper[function] = self.yc_soln.new_grid(
function.name, dimensions)
indices = [
int((i.origin if isinstance(i, LoweredDimension) else i) -
j) for i, j in zip(expr.indices, function.indices)
]
return self.mapper[function].new_relative_grid_point(indices)
elif expr.is_Add:
return nary2binary(expr.args, nfac.new_add_node)
elif expr.is_Mul:
return nary2binary(expr.args, nfac.new_multiply_node)
elif expr.is_Pow:
num, den = expr.as_numer_denom()
if num == 1:
return nfac.new_divide_node(run(num), run(den))
elif expr.is_Equality:
if expr.lhs.is_Symbol:
function = expr.lhs.base.function
assert function not in self.mapper
self.mapper[function] = run(expr.rhs)
else:
return nfac.new_equation_node(*[run(i) for i in expr.args])
else:
warning("Missing handler in Devito-YASK translation")
raise NotImplementedError
def make_yc_solution(self, namer):
"""
Create and return a YASK compiler solution object.
"""
name = namer(self.name, self.nsolutions)
yc_soln = cfac.new_solution(name)
# Redirect stdout/strerr to a string or file
if configuration.yask['dump']:
filename = 'yc_dump.%s.%s.%s.txt' % (
name, configuration['platform'], configuration['isa'])
filename = os.path.join(configuration.yask['dump'], filename)
yc_soln.set_debug_output(ofac.new_file_output(filename))
else:
yc_soln.set_debug_output(ofac.new_null_output())
# Set data type size
yc_soln.set_element_bytes(self.dtype().itemsize)
# Apply compile-time optimizations
if configuration['isa'] != 'cpp':
dimensions = [
nfac.new_domain_index(str(i)) for i in self.space_dimensions
]
# Vector folding
for i, j in zip(dimensions, configuration.yask['folding']):
yc_soln.set_fold_len(i, j)
# Unrolling
for i, j in zip(dimensions, configuration.yask['clustering']):
yc_soln.set_cluster_mult(i, j)
return yc_soln
def __init__(self, name, grid, dtype):
"""
Proxy between Devito and YASK.
A ``YaskContext`` contains N :class:`YaskKernel` and M :class:`Data`,
which have common space and time dimensions.
:param name: Unique name of the context.
:param grid: A :class:`Grid` carrying the context dimensions.
:param dtype: The data type used in kernels, as a NumPy dtype.
"""
self.name = name
self.space_dimensions = grid.dimensions
self.time_dimension = grid.stepping_dim
self.dtype = dtype
# All known solutions and grids in this context
self.solutions = []
self.grids = {}
# Build the hook kernel solution (wrapper) to create grids
yc_hook = self.make_yc_solution(namespace['jit-yc-hook'])
# Need to add dummy grids to make YASK happy
# TODO: improve me
handle = [nfac.new_domain_index(str(i)) for i in self.space_dimensions]
yc_hook.new_grid('dummy_wo_time', handle)
handle = [nfac.new_step_index(str(self.time_dimension))] + handle
yc_hook.new_grid('dummy_w_time', handle)
self.yk_hook = YaskKernel(namespace['jit-yk-hook'](name, 0), yc_hook)
def __init__(self, name, domain, dtype):
"""
Proxy between Devito and YASK.
A YaskContext contains N YaskKernel and M YaskGrids.
Solutions and grids have in common the context domain. Grids, however, may
differ in the halo region, due to a different space order. The same grid
could be used in more than one of the N solutions.
:param name: Unique name of the context.
:param domain: A mapper from space dimensions to their domain size.
:param dtype: The data type used in kernels, as a NumPy dtype.
"""
self.name = name
self.domain = domain
self.dtype = dtype
# All known solutions and grids in this context
self.solutions = []
self.grids = {}
# Build the hook kernel solution (wrapper) to create grids
yc_hook = self.make_yc_solution(namespace['jit-yc-hook'])
# Need to add dummy grids to make YASK happy
# TODO: improve me
dimensions = [nfac.new_domain_index(i) for i in domain]
yc_hook.new_grid('dummy_wo_time', dimensions)
dimensions = [nfac.new_step_index(namespace['time-dim'])] + dimensions
yc_hook.new_grid('dummy_w_time', dimensions)
self.yk_hook = YaskKernel(namespace['jit-yk-hook'](name, 0), yc_hook,
domain)
def make_yask_ast(expr, yc_soln, mapper=None):
def nary2binary(args, op):
r = make_yask_ast(args[0], yc_soln, mapper)
return r if len(args) == 1 else op(r, nary2binary(args[1:], op))
if mapper is None:
mapper = {}
if expr.is_Integer:
return nfac.new_const_number_node(int(expr))
elif expr.is_Float:
return nfac.new_const_number_node(float(expr))
elif expr.is_Rational:
a, b = expr.as_numer_denom()
return nfac.new_const_number_node(float(a) / float(b))
elif expr.is_Symbol:
function = expr.function
if function.is_Constant:
# Create a YASK var if it's the first time we encounter the embedded Function
if function not in mapper:
mapper[function] = yc_soln.new_var(function.name, [])
# Allow number of time-steps to be set in YASK kernel.
mapper[function].set_dynamic_step_alloc(True)
return mapper[function].new_var_point([])
elif function.is_Dimension:
if expr.is_Time:
return nfac.new_step_index(expr.name)
elif expr.is_Space:
# `expr.root` instead of `expr` because YASK wants the SubDimension
# information to be provided as if-conditions, and this is handled
# a-posteriori directly by `yaskit`
return nfac.new_domain_index(expr.root.name)
else:
return nfac.new_misc_index(expr.name)
else:
# E.g., A DSE-generated temporary, which must have already been
# encountered as a LHS of a previous expression
assert function in mapper
return mapper[function]
elif expr.is_Indexed:
function = expr.function
# Create a YASK var if it's the first time we encounter the embedded Function
if function not in mapper:
dimensions = [
make_yask_ast(i.root, yc_soln, mapper)
for i in function.indices
]
mapper[function] = yc_soln.new_var(function.name, dimensions)
# Allow number of time-steps to be set in YASK kernel.
mapper[function].set_dynamic_step_alloc(True)
# We also get to know some relevant Dimension-related symbols
# For example, the min point of the `x` Dimension, `x_m`, should
# be mapped to YASK's `FIRST(x)`
for d in function.indices:
node = nfac.new_domain_index(d.name)
mapper[d.symbolic_min] = nfac.new_first_domain_index(node)
mapper[d.symbolic_max] = nfac.new_last_domain_index(node)
indices = [make_yask_ast(i, yc_soln, mapper) for i in expr.indices]
return mapper[function].new_var_point(indices)
elif expr.is_Add:
return nary2binary(expr.args, nfac.new_add_node)
elif expr.is_Mul:
return nary2binary(expr.args, nfac.new_multiply_node)
elif expr.is_Pow:
base, exp = expr.as_base_exp()
if not exp.is_integer:
raise NotImplementedError("Non-integer powers unsupported in "
"Devito-YASK translation")
if int(exp) < 0:
num, den = expr.as_numer_denom()
return nfac.new_divide_node(make_yask_ast(num, yc_soln, mapper),
make_yask_ast(den, yc_soln, mapper))
elif int(exp) >= 1:
return nary2binary([base] * exp, nfac.new_multiply_node)
else:
warning("0-power found in Devito-YASK translation? setting to 1")
return nfac.new_const_number_node(1)
elif isinstance(expr, IntDiv):
return nfac.new_divide_node(make_yask_ast(expr.lhs, yc_soln, mapper),
make_yask_ast(expr.rhs, yc_soln, mapper))
elif expr.is_Equality:
if expr.lhs.is_Symbol:
function = expr.lhs.function
# The IETs are always in SSA form, so the only situation in
# which `function` may already appear in `mapper` is when we've
# already processed it as part of a different set of
# boundary conditions. For example consider `expr = a[x]*2`:
# first time, expr executed iff `x == FIRST_INDEX(x) + 7`
# second time, expr executed iff `x == FIRST_INDEX(x) + 6`
if function not in mapper:
mapper[function] = make_yask_ast(expr.rhs, yc_soln, mapper)
else:
return nfac.new_equation_node(
*[make_yask_ast(i, yc_soln, mapper) for i in expr.args])
else:
raise NotImplementedError("Missing handler in Devito-YASK translation")
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 make_yask_ast(expr, yc_soln, mapper):
def nary2binary(args, op):
r = make_yask_ast(args[0], yc_soln, mapper)
return r if len(args) == 1 else op(r, nary2binary(args[1:], op))
if expr.is_Integer:
return nfac.new_const_number_node(int(expr))
elif expr.is_Float:
return nfac.new_const_number_node(float(expr))
elif expr.is_Rational:
a, b = expr.as_numer_denom()
return nfac.new_const_number_node(float(a) / float(b))
elif expr.is_Symbol:
function = expr.function
if function.is_Constant:
# Create a YASK grid if it's the first time we encounter the embedded Function
if function not in mapper:
mapper[function] = yc_soln.new_grid(function.name, [])
# Allow number of time-steps to be set in YASK kernel.
mapper[function].set_dynamic_step_alloc(True)
return mapper[function].new_grid_point([])
elif function.is_Dimension:
if expr.is_Time:
return nfac.new_step_index(expr.name)
elif expr.is_Space:
# `expr.root` instead of `expr` because YASK wants the SubDimension
# information to be provided as if-conditions, and this is handled
# a-posteriori directly by `yaskit`
return nfac.new_domain_index(expr.root.name)
else:
return nfac.new_misc_index(expr.name)
else:
# A DSE-generated temporary, which must have already been
# encountered as a LHS of a previous expression
assert function in mapper
return mapper[function]
elif expr.is_Indexed:
function = expr.function
# Create a YASK grid if it's the first time we encounter the embedded Function
if function not in mapper:
dimensions = [
make_yask_ast(i.root, yc_soln, mapper)
for i in function.indices
]
mapper[function] = yc_soln.new_grid(function.name, dimensions)
# Allow number of time-steps to be set in YASK kernel.
mapper[function].set_dynamic_step_alloc(True)
indices = [make_yask_ast(i, yc_soln, mapper) for i in expr.indices]
return mapper[function].new_grid_point(indices)
elif expr.is_Add:
return nary2binary(expr.args, nfac.new_add_node)
elif expr.is_Mul:
return nary2binary(expr.args, nfac.new_multiply_node)
elif expr.is_Pow:
base, exp = expr.as_base_exp()
if not exp.is_integer:
raise NotImplementedError("Non-integer powers unsupported in "
"Devito-YASK translation")
if int(exp) < 0:
num, den = expr.as_numer_denom()
return nfac.new_divide_node(make_yask_ast(num, yc_soln, mapper),
make_yask_ast(den, yc_soln, mapper))
elif int(exp) >= 1:
return nary2binary([base] * exp, nfac.new_multiply_node)
else:
warning("0-power found in Devito-YASK translation? setting to 1")
return nfac.new_const_number_node(1)
elif isinstance(expr, IntDiv):
return nfac.new_divide_node(make_yask_ast(expr.lhs, yc_soln, mapper),
make_yask_ast(expr.rhs, yc_soln, mapper))
elif expr.is_Equality:
if expr.lhs.is_Symbol:
function = expr.lhs.base.function
assert function not in mapper
mapper[function] = make_yask_ast(expr.rhs, yc_soln, mapper)
else:
return nfac.new_equation_node(
*[make_yask_ast(i, yc_soln, mapper) for i in expr.args])
else:
raise NotImplementedError("Missing handler in Devito-YASK translation")
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 make_yask_ast(expr, yc_soln, mapper):
def nary2binary(args, op):
r = make_yask_ast(args[0], yc_soln, mapper)
return r if len(args) == 1 else op(r, nary2binary(args[1:], op))
if expr.is_Integer:
return nfac.new_const_number_node(int(expr))
elif expr.is_Float:
return nfac.new_const_number_node(float(expr))
elif expr.is_Rational:
a, b = expr.as_numer_denom()
return nfac.new_const_number_node(float(a)/float(b))
elif expr.is_Symbol:
function = expr.function
if function.is_Constant:
if function not in mapper:
mapper[function] = yc_soln.new_grid(function.name, [])
return mapper[function].new_grid_point([])
elif function.is_Dimension:
if expr.is_Time:
return nfac.new_step_index(expr.name)
elif expr.is_Space:
return nfac.new_domain_index(expr.name)
else:
return nfac.new_misc_index(expr.name)
else:
# A DSE-generated temporary, which must have already been
# encountered as a LHS of a previous expression
assert function in mapper
return mapper[function]
elif expr.is_Indexed:
# Create a YASK compiler grid if it's the first time we encounter a Function
function = expr.function
if function not in mapper:
dimensions = [make_yask_ast(i, yc_soln, mapper) for i in function.indices]
mapper[function] = yc_soln.new_grid(function.name, dimensions)
# Convert the Indexed into a YASK grid access
indices = []
for i in expr.indices:
if i.is_integer:
# Typically, if we end up here it's because we have a misc dimension
indices.append(make_yask_ast(i, yc_soln, mapper))
else:
# We must always use the parent ("main") dimension when creating
# YASK expressions
af = split_affine(i)
dim = af.var.parent if af.var.is_Derived else af.var
indices.append(make_yask_ast(dim + af.shift, yc_soln, mapper))
return mapper[function].new_grid_point(indices)
elif expr.is_Add:
return nary2binary(expr.args, nfac.new_add_node)
elif expr.is_Mul:
return nary2binary(expr.args, nfac.new_multiply_node)
elif expr.is_Pow:
base, exp = expr.as_base_exp()
if not exp.is_integer:
raise NotImplementedError("Non-integer powers unsupported in "
"Devito-YASK translation")
if int(exp) < 0:
num, den = expr.as_numer_denom()
return nfac.new_divide_node(make_yask_ast(num, yc_soln, mapper),
make_yask_ast(den, yc_soln, mapper))
elif int(exp) >= 1:
return nary2binary([base] * exp, nfac.new_multiply_node)
else:
warning("0-power found in Devito-YASK translation? setting to 1")
return nfac.new_const_number_node(1)
elif expr.is_Equality:
if expr.lhs.is_Symbol:
function = expr.lhs.base.function
assert function not in mapper
mapper[function] = make_yask_ast(expr.rhs, yc_soln, mapper)
else:
return nfac.new_equation_node(*[make_yask_ast(i, yc_soln, mapper)
for i in expr.args])
else:
raise NotImplementedError("Missing handler in Devito-YASK translation")
def make_yask_ast(expr, yc_soln, mapper=None):
def nary2binary(args, op):
r = make_yask_ast(args[0], yc_soln, mapper)
return r if len(args) == 1 else op(r, nary2binary(args[1:], op))
if mapper is None:
mapper = {}
if expr.is_Integer:
return nfac.new_const_number_node(int(expr))
elif expr.is_Float:
return nfac.new_const_number_node(float(expr))
elif expr.is_Rational:
a, b = expr.as_numer_denom()
return nfac.new_const_number_node(float(a)/float(b))
elif expr.is_Symbol:
function = expr.function
if function.is_Constant:
# Create a YASK grid if it's the first time we encounter the embedded Function
if function not in mapper:
mapper[function] = yc_soln.new_grid(function.name, [])
# Allow number of time-steps to be set in YASK kernel.
mapper[function].set_dynamic_step_alloc(True)
return mapper[function].new_grid_point([])
elif function.is_Dimension:
if expr.is_Time:
return nfac.new_step_index(expr.name)
elif expr.is_Space:
# `expr.root` instead of `expr` because YASK wants the SubDimension
# information to be provided as if-conditions, and this is handled
# a-posteriori directly by `yaskit`
return nfac.new_domain_index(expr.root.name)
else:
return nfac.new_misc_index(expr.name)
else:
# E.g., A DSE-generated temporary, which must have already been
# encountered as a LHS of a previous expression
assert function in mapper
return mapper[function]
elif expr.is_Indexed:
function = expr.function
# Create a YASK grid if it's the first time we encounter the embedded Function
if function not in mapper:
dimensions = [make_yask_ast(i.root, yc_soln, mapper)
for i in function.indices]
mapper[function] = yc_soln.new_grid(function.name, dimensions)
# Allow number of time-steps to be set in YASK kernel.
mapper[function].set_dynamic_step_alloc(True)
# We also get to know some relevant Dimension-related symbols
# For example, the min point of the `x` Dimension, `x_m`, should
# be mapped to YASK's `FIRST(x)`
for d in function.indices:
node = nfac.new_domain_index(d.name)
mapper[d.symbolic_min] = nfac.new_first_domain_index(node)
mapper[d.symbolic_max] = nfac.new_last_domain_index(node)
indices = [make_yask_ast(i, yc_soln, mapper) for i in expr.indices]
return mapper[function].new_grid_point(indices)
elif expr.is_Add:
return nary2binary(expr.args, nfac.new_add_node)
elif expr.is_Mul:
return nary2binary(expr.args, nfac.new_multiply_node)
elif expr.is_Pow:
base, exp = expr.as_base_exp()
if not exp.is_integer:
raise NotImplementedError("Non-integer powers unsupported in "
"Devito-YASK translation")
if int(exp) < 0:
num, den = expr.as_numer_denom()
return nfac.new_divide_node(make_yask_ast(num, yc_soln, mapper),
make_yask_ast(den, yc_soln, mapper))
elif int(exp) >= 1:
return nary2binary([base] * exp, nfac.new_multiply_node)
else:
warning("0-power found in Devito-YASK translation? setting to 1")
return nfac.new_const_number_node(1)
elif isinstance(expr, IntDiv):
return nfac.new_divide_node(make_yask_ast(expr.lhs, yc_soln, mapper),
make_yask_ast(expr.rhs, yc_soln, mapper))
elif expr.is_Equality:
if expr.lhs.is_Symbol:
function = expr.lhs.function
# The IETs are always in SSA form, so the only situation in
# which `function` may already appear in `mapper` is when we've
# already processed it as part of a different set of
# boundary conditions. For example consider `expr = a[x]*2`:
# first time, expr executed iff `x == FIRST_INDEX(x) + 7`
# second time, expr executed iff `x == FIRST_INDEX(x) + 6`
if function not in mapper:
mapper[function] = make_yask_ast(expr.rhs, yc_soln, mapper)
else:
return nfac.new_equation_node(*[make_yask_ast(i, yc_soln, mapper)
for i in expr.args])
else:
raise NotImplementedError("Missing handler in Devito-YASK translation")
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
