def make_grid(self, obj):
"""
Create and return a new :class:`Data`, a YASK grid wrapper. Memory
is allocated.
:param obj: The :class:`Function` for which a YASK grid is allocated.
"""
if set(obj.indices) < set(self.space_dimensions):
exit("Need a Function[x,y,z] to create a YASK grid.")
name = 'devito_%s_%d' % (obj.name, contexts.ngrids)
# Set up the YASK grid and allocate memory
grid = self.yk_hook.new_grid(name, obj)
for i, s, h in zip(obj.indices, obj.shape, obj._halo):
if i.is_Time:
assert grid.is_dim_used(i.name)
assert grid.get_alloc_size(i.name) == s
else:
# Note:
# 1) The halo is set to a value which is the max between the number
# of points on the left and the number of points on the right of
# the approximation (the same with a centered approximation)
# 2) from the YASK docs: "If the halo is set to a value larger than
# the padding size, the padding size will be automatically increased
# to accomodate it
grid.set_halo_size(i.name, max(h))
grid.alloc_storage()
self.grids[name] = grid
return grid
def find_offloadable_trees(nodes):
"""
Return the trees within ``nodes`` that can be computed by YASK.
A tree is "offloadable to YASK" if it is embedded in a time stepping loop
*and* all of the grids accessed by the enclosed equations are homogeneous
(i.e., same dimensions and data type).
"""
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.grid, i.dtype) for i in functions if i.is_TimeFunction)
if len(keys) == 0:
continue
elif len(keys) > 1:
exit("Cannot handle Operators w/ heterogeneous grids")
grid, dtype = keys.pop()
# Is this a "complete" tree iterating over the entire grid?
dims = [i.dim for i in tree]
if all(i in dims for i in grid.dimensions) and\
any(j in dims for j in [grid.time_dim, grid.stepping_dim]):
offloadable.append((tree, grid, dtype))
return offloadable
def find_offloadable_trees(nodes):
"""
Return the trees within ``nodes`` that can be computed by YASK.
A tree is "offloadable to YASK" if it is embedded in a time stepping loop
*and* all of the grids accessed by the enclosed equations are homogeneous
(i.e., same dimensions, shape, data type).
"""
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_TimeFunction)
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))
return offloadable
def fetch(self, dimensions, shape, dtype):
"""
Fetch the :class:`YaskContext` in ``self`` uniquely identified by
``dimensions``, ``shape``, and ``dtype``. Create a new (empty)
:class:`YaskContext` on miss.
"""
# Sanity checks
assert len(dimensions) == len(shape)
dimensions = [str(i) for i in dimensions]
if set(dimensions) < {'x', 'y', 'z'}:
exit("Need a Function[x,y,z] for initialization")
# The time dimension is dropped as implicit to the context
domain = OrderedDict([(i, j) for i, j in zip(dimensions, shape)
if i != namespace['time-dim']])
# A unique key for this context.
key = tuple([configuration['isa'], dtype] + list(domain.items()))
# Fetch or create a YaskContext
if key in self:
log("Fetched existing context from cache")
else:
self[key] = YaskContext('ctx%d' % self.ncontexts, domain, dtype)
self.ncontexts += 1
log("Context successfully created!")
return self[key]
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 fetch(self, dimensions, dtype):
"""
Fetch the :class:`YaskContext` in ``self`` uniquely identified by
``dimensions`` and ``dtype``.
"""
key = self._getkey(None, dtype, dimensions)
context = self.get(key, self._partial_map.get(key))
if context is not None:
log("Fetched existing YaskContext from cache")
return context
else:
exit("Couldn't find YaskContext for key=`%s`" % str(key))
def fetch(self, dimensions, dtype):
"""
Fetch the YaskContext in ``self`` uniquely identified by ``dimensions`` and
``dtype``.
"""
key = self._getkey(None, dtype, dimensions)
context = self.get(key, self._partial_map.get(key))
if context is not None:
log("Fetched existing YaskContext from cache")
return context
else:
exit("Couldn't find YaskContext for key=`%s`" % str(key))
def fetch(self, grid, dtype, dimensions=None):
"""
Fetch the :class:`YaskContext` in ``self`` uniquely identified by
``grid`` and ``dtype``.
"""
key = self._getkey(grid, dtype, dimensions)
context = self.get(key, self._partial_map.get(key))
if context is not None:
log("Fetched existing context from cache")
return context
else:
exit("Couldn't find context for grid %s" % grid)
def make_grid(self, obj):
"""
Create and return a new :class:`YaskGrid`, a YASK grid wrapper. Memory
is allocated.
:param obj: The symbolic data object for which a YASK grid is allocated.
"""
if set(obj.indices) < set(self.space_dimensions):
exit("Need a Function[x,y,z] to create a YASK grid.")
name = 'devito_%s_%d' % (obj.name, contexts.ngrids)
log("Allocating YaskGrid for %s (%s)" % (obj.name, str(obj.shape)))
grid = self.yk_hook.new_grid(name, obj)
wrapper = YaskGrid(grid, obj.shape, obj.space_order, obj.dtype)
self.grids[name] = wrapper
return wrapper
def make_grid(self, obj):
"""
Create and return a new :class:`Data`, a YASK grid wrapper. Memory
is allocated.
:param obj: The :class:`Function` for which a YASK grid is allocated.
"""
if set(obj.indices) < set(self.space_dimensions):
exit("Need a Function[x,y,z] to create a YASK grid.")
name = 'devito_%s_%d' % (obj.name, contexts.ngrids)
# Create the YASK grid
grid = self.yk_hook.new_grid(name, obj)
# Where should memory be allocated ?
alloc = obj._allocator
if alloc.is_Numa:
if alloc.put_onnode:
grid.set_numa_preferred(alloc.node)
elif alloc.put_local:
grid.set_numa_preferred(namespace['numa-put-local'])
for i, s, h in zip(obj.indices, obj.shape_allocated, obj._extent_halo):
if i.is_Time:
assert grid.is_dim_used(i.name)
assert grid.get_alloc_size(i.name) == s
else:
# Note:
# From the YASK docs: "If the halo is set to a value larger than
# the padding size, the padding size will be automatically increased
# to accomodate it."
grid.set_left_halo_size(i.name, h.left)
grid.set_right_halo_size(i.name, h.right)
grid.alloc_storage()
self.grids[name] = grid
return grid
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 __init__(self, name, yc_soln, local_grids=None):
"""
Write out a YASK kernel, build it using YASK's Makefiles,
import the corresponding SWIG-generated Python module, and finally
create a YASK kernel solution object.
:param name: Unique name of this YaskKernel.
:param yc_soln: YaskCompiler solution.
:param local_grids: A local grid is necessary to run the YaskKernel,
but its final content can be ditched. Indeed, local
grids are hidden to users -- for example, they could
represent temporary arrays introduced by the DSE.
This parameter tells which of the ``yc_soln``'s grids
are local.
"""
self.name = name
# Shared object name
self.soname = "%s.%s.%s" % (name, yc_soln.get_name(),
configuration['platform'])
# It's necessary to `clean` the YASK kernel directory *before*
# writing out the first `yask_stencil_code.hpp`
make(namespace['path'], ['-C', namespace['kernel-path'], 'clean'])
# Write out the stencil file
if not os.path.exists(namespace['kernel-path-gen']):
os.makedirs(namespace['kernel-path-gen'])
yc_soln.format(configuration['isa'],
ofac.new_file_output(namespace['kernel-output']))
# JIT-compile it
try:
compiler = configuration.yask['compiler']
opt_level = 1 if configuration.yask['develop-mode'] else 3
make(
namespace['path'],
[
'-j3',
'YK_CXX=%s' % compiler.cc,
'YK_CXXOPT=-O%d' % opt_level,
'mpi=0', # Disable MPI for now
# "EXTRA_MACROS=TRACE",
'YK_BASE=%s' % str(name),
'stencil=%s' % yc_soln.get_name(),
'arch=%s' % configuration['platform'],
'-C',
namespace['kernel-path'],
'api'
])
except CompilationError:
exit("Kernel solution compilation")
# Import the corresponding Python (SWIG-generated) module
try:
yk = getattr(__import__('yask', fromlist=[name]), name)
except ImportError:
exit("Python YASK kernel bindings")
try:
yk = reload(yk)
except NameError:
# Python 3.5 compatibility
yk = importlib.reload(yk)
# Create the YASK solution object
kfac = yk.yk_factory()
self.env = kfac.new_env()
self.soln = kfac.new_solution(self.env)
# MPI setup: simple rank configuration in 1st dim only.
# TODO: in production runs, the ranks would be distributed along all
# domain dimensions.
self.soln.set_num_ranks(self.space_dimensions[0],
self.env.get_num_ranks())
# Redirect stdout/strerr to a string or file
if configuration.yask['dump']:
filename = 'yk_dump.%s.%s.%s.txt' % (
self.name, configuration['platform'], configuration['isa'])
filename = os.path.join(configuration.yask['dump'], filename)
self.output = yk.yask_output_factory().new_file_output(filename)
else:
self.output = yk.yask_output_factory().new_string_output()
self.soln.set_debug_output(self.output)
# Users may want to run the same Operator (same domain etc.) with
# different grids.
self.grids = {i.get_name(): i for i in self.soln.get_grids()}
self.local_grids = {
i.name: self.grids[i.name]
for i in (local_grids or [])
}
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 run_py(self, ntimesteps):
exit("Cannot run a NullSolution through YASK's Python bindings")
def __init__(self, name, yc_soln, domain):
"""
Write out a YASK kernel, build it using YASK's Makefiles,
import the corresponding SWIG-generated Python module, and finally
create a YASK kernel solution object.
:param name: Unique name of this YaskKernel.
:param yc_soln: YaskCompiler solution.
:param domain: A mapper from space dimensions to their domain size.
"""
self.name = name
# Shared object name
self.soname = "%s.%s.%s" % (name, yc_soln.get_name(),
yask_configuration['arch'])
# It's necessary to `clean` the YASK kernel directory *before*
# writing out the first `yask_stencil_code.hpp`
make(namespace['path'], ['-C', namespace['kernel-path'], 'clean'])
# Write out the stencil file
if not os.path.exists(namespace['kernel-path-gen']):
os.makedirs(namespace['kernel-path-gen'])
yc_soln.format(yask_configuration['isa'],
ofac.new_file_output(namespace['kernel-output']))
# JIT-compile it
try:
opt_level = 1 if yask_configuration['develop-mode'] else 3
make(
os.environ['YASK_HOME'],
[
'-j',
'YK_CXXOPT=-O%d' % opt_level,
# "EXTRA_MACROS=TRACE",
'YK_BASE=%s' % str(name),
'stencil=%s' % yc_soln.get_name(),
'arch=%s' % yask_configuration['arch'],
'-C',
namespace['kernel-path'],
'api'
])
except CompilationError:
exit("Kernel solution compilation")
# Import the corresponding Python (SWIG-generated) module
try:
yk = importlib.import_module(name)
except ImportError:
exit("Python YASK kernel bindings")
try:
yk = reload(yk)
except NameError:
# Python 3.5 compatibility
yk = importlib.reload(yk)
# Create the YASK solution object
kfac = yk.yk_factory()
self.env = kfac.new_env()
self.soln = kfac.new_solution(self.env)
# MPI setup: simple rank configuration in 1st dim only.
# TODO: in production runs, the ranks would be distributed along all
# domain dimensions.
self.soln.set_num_ranks(self.soln.get_domain_dim_names()[0],
self.env.get_num_ranks())
# Redirect stdout/strerr to a string
self.output = yk.yask_output_factory().new_string_output()
self.soln.set_debug_output(self.output)
# Set up the solution domain size
for k, v in domain.items():
self.soln.set_rank_domain_size(k, v)
def __init__(self, name, yc_soln, local_grids=None):
"""
Write out a YASK kernel, build it using YASK's Makefiles,
import the corresponding SWIG-generated Python module, and finally
create a YASK kernel solution object.
:param name: Unique name of this YaskKernel.
:param yc_soln: YaskCompiler solution.
:param local_grids: A local grid is necessary to run the YaskKernel,
but its final content can be ditched. Indeed, local
grids are hidden to users -- for example, they could
represent temporary arrays introduced by the DSE.
This parameter tells which of the ``yc_soln``'s grids
are local.
"""
self.name = name
# Shared object name
self.soname = "%s.devito.%s" % (name, configuration['platform'])
# The directory in which the YASK-generated code (.hpp) will be placed
yk_codegen = namespace['yask-codegen'](name, 'devito',
configuration['platform'])
if not os.path.exists(yk_codegen):
os.makedirs(yk_codegen)
# Write out the stencil file
yk_codegen_file = os.path.join(yk_codegen,
namespace['yask-codegen-file'])
yc_soln.format(configuration['isa'],
ofac.new_file_output(yk_codegen_file))
# JIT-compile it
try:
compiler = configuration.yask['compiler']
if configuration['develop-mode']:
if yc_soln.get_num_equations() == 0:
# YASK will compile more quickly, and no price has to be paid
# in terms of performance, as this is a void kernel
opt_level = 0
else:
opt_level = 1
else:
opt_level = 3
args = [
'-j',
'YK_CXX=%s' % compiler.cc,
'YK_CXXOPT=-O%d' % opt_level,
# No MPI support at the moment
'mpi=0',
# To locate the YASK compiler
'YC_EXEC=%s' % os.path.join(namespace['path'], 'bin'),
# Error out if a grid not explicitly defined in the compiler is created
'allow_new_grid_types=0',
# To give a unique name to the generated Python modules, rather
# than creating `yask_kernel.py`
'YK_BASE=%s' % name,
# `stencil` and `arch` should always be provided
'stencil=%s' % 'devito',
'arch=%s' % configuration['platform'],
# The root directory of generated code files, shared libs, Python modules
'YASK_OUTPUT_DIR=%s' % namespace['yask-output-dir'],
# Pick the YASK kernel Makefile, i.e. the one under `yask/src/kernel`
'-C',
namespace['kernel-path'],
'api'
]
# Other potentially useful args:
# - "EXTRA_MACROS=TRACE", -- debugging option
make(namespace['path'], args)
except CompilationError:
exit("Kernel solution compilation")
# Import the corresponding Python (SWIG-generated) module
try:
sys.path.append(os.path.join(namespace['yask-output-dir'], 'yask'))
importlib.invalidate_caches()
yk = importlib.import_module(name)
except ImportError:
exit("Python YASK kernel bindings")
# Create the YASK solution object
kfac = yk.yk_factory()
self.env = kfac.new_env()
self.soln = kfac.new_solution(self.env)
# Apply any user-provided options, if any.
# These are applied here instead of just before prepare_solution()
# so that applicable options will apply to all API calls.
self.soln.apply_command_line_options(configuration.yask['options']
or '')
# MPI setup: simple rank configuration in 1st dim only.
# TODO: in production runs, the ranks would be distributed along all
# domain dimensions.
self.soln.set_num_ranks(self.space_dimensions[0],
self.env.get_num_ranks())
# Redirect stdout to a string or file
if configuration.yask['dump']:
filename = 'yk_dump.%s.%s.%s.txt' % (
name, configuration['platform'], configuration['isa'])
filename = os.path.join(configuration.yask['dump'], filename)
self.output = yk.yask_output_factory().new_file_output(filename)
else:
self.output = yk.yask_output_factory().new_string_output()
self.soln.set_debug_output(self.output)
# Users may want to run the same Operator (same domain etc.) with
# different grids.
self.grids = {i.get_name(): i for i in self.soln.get_grids()}
self.local_grids = {
i.name: self.grids[i.name]
for i in (local_grids or [])
}
