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 __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 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_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 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 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")