def _specialize_iet(self, nodes):
"""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."""
log("Specializing a Devito Operator for YASK...")
self.context = YaskNullContext()
self.yk_soln = YaskNullKernel()
offloadable = find_offloadable_trees(nodes)
if len(offloadable) == 0:
log("No offloadable trees found")
elif len(offloadable) == 1:
tree, grid, dtype = offloadable[0]
self.context = contexts.fetch(grid, dtype)
# Create a YASK compiler solution for this Operator
yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln'])
transform = sympy2yask(self.context, yc_soln)
try:
for i in tree[-1].nodes:
transform(i.expr)
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
nodes = Transformer({tree[1]: funcall}).visit(nodes)
# Track /funcall/ as an external function call
self.func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
local_grids = [i for i in transform.mapper if i.is_Array]
self.yk_soln = self.context.make_yk_solution(
namespace['jit-yk-soln'], yc_soln, local_grids)
# 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:
log("Unable to offload a candidate tree.")
else:
exit("Found more than one offloadable trees in a single Operator")
# 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
nodes = make_grid_accesses(nodes)
log("Specialization successfully performed!")
return nodes
def _specialize(self, nodes, parameters):
"""
Create a YASK representation of this Iteration/Expression tree.
``parameters`` is modified in-place adding YASK-related arguments.
"""
log("Specializing a Devito Operator for YASK...")
self.context = YaskNullContext()
self.yk_soln = YaskNullKernel()
local_grids = []
offloadable = find_offloadable_trees(nodes)
if len(offloadable) == 0:
log("No offloadable trees found")
elif len(offloadable) == 1:
tree, grid, dtype = offloadable[0]
self.context = contexts.fetch(grid, dtype)
# Create a YASK compiler solution for this Operator
yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln'])
transform = sympy2yask(self.context, yc_soln)
try:
for i in tree[-1].nodes:
transform(i.expr)
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
nodes = Transformer({tree[1]: funcall}).visit(nodes)
# Track /funcall/ as an external function call
self.func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
local_grids += [i for i in transform.mapper if i.is_Array]
self.yk_soln = self.context.make_yk_solution(
namespace['jit-yk-soln'], yc_soln, local_grids)
# Now we must drop a pointer to the YASK solution down to C-land
parameters.append(
Object(namespace['code-soln-name'],
namespace['type-solution'],
self.yk_soln.rawpointer))
# 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:
log("Unable to offload a candidate tree.")
else:
exit("Found more than one offloadable trees in a single Operator")
# 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
nodes = make_grid_accesses(nodes)
# Update the parameters list adding all necessary YASK grids
for i in list(parameters) + local_grids:
try:
if i.from_YASK:
parameters.append(
Object(namespace['code-grid-name'](i.name),
namespace['type-grid']))
except AttributeError:
# Ignore e.g. Dimensions
pass
log("Specialization successfully performed!")
return nodes
class Operator(OperatorRunnable):
"""
A special :class:`OperatorCore` to JIT-compile and run operators through YASK.
"""
_default_headers = OperatorRunnable._default_headers
_default_headers += ['#define restrict __restrict']
_default_includes = OperatorRunnable._default_includes + [
'yask_kernel_api.hpp'
]
def __init__(self, expressions, **kwargs):
kwargs['dle'] = ('denormals', ) + (
('openmp', ) if configuration['openmp'] else ())
super(Operator, self).__init__(expressions, **kwargs)
# Each YASK Operator needs to have its own compiler (hence the copy()
# below) because Operator-specific shared object will be added to the
# list of linked libraries
self._compiler = configuration.yask['compiler'].copy()
def _specialize(self, nodes, parameters):
"""
Create a YASK representation of this Iteration/Expression tree.
``parameters`` is modified in-place adding YASK-related arguments.
"""
log("Specializing a Devito Operator for YASK...")
self.context = YaskNullContext()
self.yk_soln = YaskNullKernel()
local_grids = []
offloadable = find_offloadable_trees(nodes)
if len(offloadable) == 0:
log("No offloadable trees found")
elif len(offloadable) == 1:
tree, grid, dtype = offloadable[0]
self.context = contexts.fetch(grid, dtype)
# Create a YASK compiler solution for this Operator
yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln'])
transform = sympy2yask(self.context, yc_soln)
try:
for i in tree[-1].nodes:
transform(i.expr)
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
nodes = Transformer({tree[1]: funcall}).visit(nodes)
# Track /funcall/ as an external function call
self.func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
local_grids += [i for i in transform.mapper if i.is_Array]
self.yk_soln = self.context.make_yk_solution(
namespace['jit-yk-soln'], yc_soln, local_grids)
# Now we must drop a pointer to the YASK solution down to C-land
parameters.append(
Object(namespace['code-soln-name'],
namespace['type-solution'],
self.yk_soln.rawpointer))
# 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:
log("Unable to offload a candidate tree.")
else:
exit("Found more than one offloadable trees in a single Operator")
# 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
nodes = make_grid_accesses(nodes)
# Update the parameters list adding all necessary YASK grids
for i in list(parameters) + local_grids:
try:
if i.from_YASK:
parameters.append(
Object(namespace['code-grid-name'](i.name),
namespace['type-grid']))
except AttributeError:
# Ignore e.g. Dimensions
pass
log("Specialization successfully performed!")
return nodes
def arguments(self, **kwargs):
mapper = {i.name: i for i in self.parameters}
local_grids_mapper = {
namespace['code-grid-name'](k): v
for k, v in self.yk_soln.local_grids.items()
}
# The user has the illusion to provide plain data objects to the
# generated kernels, but what we actually need and thus going to
# provide are pointers to the wrapped YASK grids.
toshare = {}
more_args = {}
for i in self.parameters:
grid_arg = mapper.get(namespace['code-grid-name'](i.name))
if grid_arg is not None:
assert i.provider.from_YASK is True
obj = kwargs.get(i.name, i.provider)
# Get the associated Data wrapper (scalars are a special case)
if np.isscalar(obj):
wrapper = DataScalar(obj)
toshare[i.provider] = wrapper
else:
wrapper = obj.data
toshare[obj] = wrapper
# Add C-level pointer to the YASK grids
more_args[grid_arg.name] = wrapper.rawpointer
elif i.name in local_grids_mapper:
# Add C-level pointer to the temporary YASK grids
more_args[i.name] = rawpointer(local_grids_mapper[i.name])
kwargs.update(more_args)
return super(Operator, self).arguments(**kwargs), toshare
def apply(self, **kwargs):
# Build the arguments list to invoke the kernel function
arguments, toshare = self.arguments(**kwargs)
log("Running YASK Operator through Devito...")
self.yk_soln.run(self.cfunction, arguments, toshare)
log("YASK Operator successfully run!")
# Output summary of performance achieved
return self._profile_output(arguments)
@property
def compile(self):
"""
JIT-compile the C code generated by the Operator.
It is ensured that JIT compilation will only be performed once per
:class:`Operator`, reagardless of how many times this method is invoked.
:returns: The file name of the JIT-compiled function.
"""
if self._lib is None:
# No need to recompile if a shared object has already been loaded.
if not isinstance(self.yk_soln, YaskNullKernel):
self._compiler.libraries.append(self.yk_soln.soname)
return jit_compile(self.ccode, self._compiler)
else:
return self._lib.name
class Operator(OperatorRunnable):
"""
A special :class:`OperatorCore` to JIT-compile and run operators through YASK.
"""
_default_headers = OperatorRunnable._default_headers
_default_headers += ['#define restrict __restrict']
_default_includes = OperatorRunnable._default_includes + [
'yask_kernel_api.hpp'
]
def __init__(self, expressions, **kwargs):
kwargs['dle'] = ('denormals', ) + (
('openmp', ) if configuration['openmp'] else ())
super(Operator, self).__init__(expressions, **kwargs)
# Each YASK Operator needs to have its own compiler (hence the copy()
# below) because Operator-specific shared object will be added to the
# list of linked libraries
self._compiler = configuration.yask['compiler'].copy()
def _specialize(self, nodes):
"""Create a YASK representation of this Iteration/Expression tree."""
log("Specializing a Devito Operator for YASK...")
self.context = YaskNullContext()
self.yk_soln = YaskNullKernel()
offloadable = find_offloadable_trees(nodes)
if len(offloadable) == 0:
log("No offloadable trees found")
elif len(offloadable) == 1:
tree, grid, dtype = offloadable[0]
self.context = contexts.fetch(grid, dtype)
# Create a YASK compiler solution for this Operator
yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln'])
transform = sympy2yask(self.context, yc_soln)
try:
for i in tree[-1].nodes:
transform(i.expr)
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
nodes = Transformer({tree[1]: funcall}).visit(nodes)
# Track /funcall/ as an external function call
self.func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
local_grids = [i for i in transform.mapper if i.is_Array]
self.yk_soln = self.context.make_yk_solution(
namespace['jit-yk-soln'], yc_soln, local_grids)
# 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:
log("Unable to offload a candidate tree.")
else:
exit("Found more than one offloadable trees in a single Operator")
# 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
nodes = make_grid_accesses(nodes)
log("Specialization successfully performed!")
return nodes
def _build_parameters(self, nodes):
parameters = super(Operator, self)._build_parameters(nodes)
# Add parameters "disappeared" due to offloading
parameters += tuple(i for i in self.input if i not in parameters)
return parameters
def _build_casts(self, nodes):
nodes = super(Operator, self)._build_casts(nodes)
# Add YASK solution pointer for use in C-land
soln_obj = Object(namespace['code-soln-name'],
namespace['type-solution'])
# Add YASK user and local grids pointers for use in C-land
grid_objs = [YaskGridObject(i.name) for i in self.input if i.from_YASK]
grid_objs.extend([YaskGridObject(i) for i in self.yk_soln.local_grids])
# Build pointer casts
casts = [PointerCast(soln_obj)] + [PointerCast(i) for i in grid_objs]
return List(body=casts + [nodes])
@cached_property
def _argument_defaults(self):
default_args = super(Operator, self)._argument_defaults
# Add in solution pointer
default_args[namespace['code-soln-name']] = self.yk_soln.rawpointer
# Add in local grids pointers
for k, v in self.yk_soln.local_grids.items():
default_args[namespace['code-grid-name'](k)] = rawpointer(v)
return default_args
def apply(self, **kwargs):
# Build the arguments list to invoke the kernel function
arguments = self.arguments(**kwargs)
# Map default Functions to runtime Functions; will be used for "grid sharing"
toshare = {}
for i in self.input:
v = kwargs.get(i.name, i)
if np.isscalar(v):
toshare[i] = DataScalar(v)
elif i.from_YASK and (i.is_Constant or i.is_Function):
toshare[v] = v.data
log("Running YASK Operator through Devito...")
arg_values = [arguments[p.name] for p in self.parameters]
self.yk_soln.run(self.cfunction, arg_values, toshare)
log("YASK Operator successfully run!")
# Output summary of performance achieved
return self._profile_output(arguments)
@property
def compile(self):
"""
JIT-compile the C code generated by the Operator.
It is ensured that JIT compilation will only be performed once per
:class:`Operator`, reagardless of how many times this method is invoked.
:returns: The file name of the JIT-compiled function.
"""
if self._lib is None:
# No need to recompile if a shared object has already been loaded.
if not isinstance(self.yk_soln, YaskNullKernel):
self._compiler.libraries.append(self.yk_soln.soname)
return jit_compile(self.ccode, self._compiler)
else:
return self._lib.name
def _specialize(self, nodes, parameters):
"""
Create a YASK representation of this Iteration/Expression tree.
``parameters`` is modified in-place adding YASK-related arguments.
"""
log("Specializing a Devito Operator for YASK...")
# Find offloadable Iteration/Expression trees
offloadable = []
for tree in retrieve_iteration_tree(nodes):
parallel = filter_iterations(tree, lambda i: i.is_Parallel)
if not parallel:
# Cannot offload non-parallel loops
continue
if not (IsPerfectIteration().visit(tree)
and all(i.is_Expression for i in tree[-1].nodes)):
# Don't know how to offload this Iteration/Expression to YASK
continue
functions = flatten(i.functions for i in tree[-1].nodes)
keys = set((i.indices, i.shape, i.dtype) for i in functions
if i.is_TimeData)
if len(keys) == 0:
continue
elif len(keys) > 1:
exit("Cannot handle Operators w/ heterogeneous grids")
dimensions, shape, dtype = keys.pop()
if len(dimensions) == len(tree) and\
all(i.dim == j for i, j in zip(tree, dimensions)):
# Detected a "full" Iteration/Expression tree (over both
# time and space dimensions)
offloadable.append((tree, dimensions, shape, dtype))
# Construct YASK ASTs given Devito expressions. New grids may be allocated.
if len(offloadable) == 0:
# No offloadable trees found
self.context = YaskNullContext()
self.yk_soln = YaskNullSolution()
processed = nodes
log("No offloadable trees found")
elif len(offloadable) == 1:
# Found *the* offloadable tree for this Operator
tree, dimensions, shape, dtype = offloadable[0]
self.context = contexts.fetch(dimensions, shape, dtype)
# Create a YASK compiler solution for this Operator
# Note: this can be dropped as soon as the kernel has been built
yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln'])
transform = sympy2yask(self.context, yc_soln)
try:
for i in tree[-1].nodes:
transform(i.expr)
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
processed = Transformer({tree[1]: funcall}).visit(nodes)
# Track this is an external function call
self.func_table[namespace['code-soln-run']] = FunMeta(
None, False)
# JIT-compile the newly-created YASK kernel
self.yk_soln = self.context.make_yk_solution(
namespace['jit-yk-soln'], yc_soln)
# Now we must drop a pointer to the YASK solution down to C-land
parameters.append(
Object(namespace['code-soln-name'],
namespace['type-solution'],
self.yk_soln.rawpointer))
# 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:
self.yk_soln = YaskNullSolution()
processed = nodes
log("Unable to offload a candidate tree.")
else:
exit("Found more than one offloadable trees in a single Operator")
# 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
processed = make_grid_accesses(processed)
# Update the parameters list adding all necessary YASK grids
for i in list(parameters):
try:
if i.from_YASK:
parameters.append(
Object(namespace['code-grid-name'](i.name),
namespace['type-grid'], i.data.rawpointer))
except AttributeError:
# Ignore e.g. Dimensions
pass
log("Specialization successfully performed!")
return processed
class Operator(OperatorRunnable):
"""
A special :class:`OperatorCore` to JIT-compile and run operators through YASK.
"""
_default_headers = OperatorRunnable._default_headers
_default_headers += ['#define restrict __restrict']
_default_includes = OperatorRunnable._default_includes + [
'yask_kernel_api.hpp'
]
def __init__(self, expressions, **kwargs):
kwargs['dle'] = 'basic'
super(Operator, self).__init__(expressions, **kwargs)
# Each YASK Operator needs to have its own compiler (hence the copy()
# below) because Operator-specific shared object will be added to the
# list of linked libraries
self._compiler = yask_configuration['compiler'].copy()
def _specialize(self, nodes, parameters):
"""
Create a YASK representation of this Iteration/Expression tree.
``parameters`` is modified in-place adding YASK-related arguments.
"""
log("Specializing a Devito Operator for YASK...")
# Find offloadable Iteration/Expression trees
offloadable = []
for tree in retrieve_iteration_tree(nodes):
parallel = filter_iterations(tree, lambda i: i.is_Parallel)
if not parallel:
# Cannot offload non-parallel loops
continue
if not (IsPerfectIteration().visit(tree)
and all(i.is_Expression for i in tree[-1].nodes)):
# Don't know how to offload this Iteration/Expression to YASK
continue
functions = flatten(i.functions for i in tree[-1].nodes)
keys = set((i.indices, i.shape, i.dtype) for i in functions
if i.is_TimeData)
if len(keys) == 0:
continue
elif len(keys) > 1:
exit("Cannot handle Operators w/ heterogeneous grids")
dimensions, shape, dtype = keys.pop()
if len(dimensions) == len(tree) and\
all(i.dim == j for i, j in zip(tree, dimensions)):
# Detected a "full" Iteration/Expression tree (over both
# time and space dimensions)
offloadable.append((tree, dimensions, shape, dtype))
# Construct YASK ASTs given Devito expressions. New grids may be allocated.
if len(offloadable) == 0:
# No offloadable trees found
self.context = YaskNullContext()
self.yk_soln = YaskNullSolution()
processed = nodes
log("No offloadable trees found")
elif len(offloadable) == 1:
# Found *the* offloadable tree for this Operator
tree, dimensions, shape, dtype = offloadable[0]
self.context = contexts.fetch(dimensions, shape, dtype)
# Create a YASK compiler solution for this Operator
# Note: this can be dropped as soon as the kernel has been built
yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln'])
transform = sympy2yask(self.context, yc_soln)
try:
for i in tree[-1].nodes:
transform(i.expr)
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
processed = Transformer({tree[1]: funcall}).visit(nodes)
# Track this is an external function call
self.func_table[namespace['code-soln-run']] = FunMeta(
None, False)
# JIT-compile the newly-created YASK kernel
self.yk_soln = self.context.make_yk_solution(
namespace['jit-yk-soln'], yc_soln)
# Now we must drop a pointer to the YASK solution down to C-land
parameters.append(
Object(namespace['code-soln-name'],
namespace['type-solution'],
self.yk_soln.rawpointer))
# 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:
self.yk_soln = YaskNullSolution()
processed = nodes
log("Unable to offload a candidate tree.")
else:
exit("Found more than one offloadable trees in a single Operator")
# 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
processed = make_grid_accesses(processed)
# Update the parameters list adding all necessary YASK grids
for i in list(parameters):
try:
if i.from_YASK:
parameters.append(
Object(namespace['code-grid-name'](i.name),
namespace['type-grid'], i.data.rawpointer))
except AttributeError:
# Ignore e.g. Dimensions
pass
log("Specialization successfully performed!")
return processed
def arguments(self, **kwargs):
# The user has the illusion to provide plain data objects to the
# generated kernels, but what we actually need and thus going to
# provide are pointers to the wrapped YASK grids
shadow_kwargs = {}
for k, v in kwargs.items():
try:
if v.from_YASK is True:
shadow_kwargs[namespace['code-grid-name'](k)] = v
except AttributeError:
pass
for i in self.parameters:
# Assign the yk::yk_grid pointers
obj = shadow_kwargs.get(i.name)
if i.is_PtrArgument and obj is not None:
assert (i.verify(obj.data.rawpointer))
# Also need to update this Operator's grids by sharing user-provided data
for i in self.parameters:
obj = kwargs.get(i.name)
if obj is not None and obj.from_YASK:
kgrid = self.yk_soln.grids.get(i.name)
if kgrid is not None:
obj.data.give_storage(kgrid)
else:
# A YaskGrid read/written by this Operator but not appearing
# in the offloaded YASK kernel (if any)
pass
return super(Operator, self).arguments(**kwargs)
def apply(self, **kwargs):
# Build the arguments list to invoke the kernel function
arguments, dim_sizes = self.arguments(**kwargs)
# Print some info about the solution.
log("Stencil-solution '%s':" % self.yk_soln.name)
log(" Step dimension: %s" % self.context.time_dimension)
log(" Domain dimensions: %s" % str(self.context.space_dimensions))
log(" Grids:")
for grid in self.yk_soln.grids.values():
pad = str(
[grid.get_pad_size(i) for i in self.context.space_dimensions])
log(" %s%s, pad=%s" %
(grid.get_name(), str(grid.get_dim_names()), pad))
# Required by YASK before running any stencils
self.yk_soln.prepare()
if yask_configuration['python-exec']:
log("Running YASK Operator through YASK...")
self.yk_soln.run(dim_sizes[self.context.time_dimension])
else:
log("Running YASK Operator through Devito...")
self.cfunction(*list(arguments.values()))
log("YASK Operator successfully run!")
# Output summary of performance achieved
return self._profile_output(dim_sizes)
@property
def compile(self):
"""
JIT-compile the C code generated by the Operator.
It is ensured that JIT compilation will only be performed once per
:class:`Operator`, reagardless of how many times this method is invoked.
:returns: The file name of the JIT-compiled function.
"""
if self._lib is None:
# No need to recompile if a shared object has already been loaded.
if not isinstance(self.yk_soln, YaskNullSolution):
self._compiler.libraries.append(self.yk_soln.soname)
return jit_compile(self.ccode, self._compiler)
else:
return self._lib.name