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.
"""
# 'hook' compiler solution: describes the grid
# 'hook' kernel solution: allocates memory
# A unique name for the 'hook' compiler and kernel solutions
suffix = Signer._digest(self, obj, configuration,
YaskContext._hookcounter)
YaskContext._hookcounter += 1
name = namespace['jit-hook'](suffix)
# Create 'hook' compiler solution
yc_hook = self.make_yc_solution(name)
if obj.indices != self.dimensions:
# Note: YASK wants *at least* a grid with *all* space (domain) dimensions
# *and* the stepping dimension. `obj`, however, may actually employ a
# different set of dimensions (e.g., a strict subset and/or some misc
# dimensions). In such a case, an extra dummy grid is attached
# `obj` examples: u(x, d), u(x, y, z)
dimensions = [
make_yask_ast(i, yc_hook, {}) for i in self.dimensions
]
yc_hook.new_grid('dummy_grid_full', dimensions)
dimensions = [make_yask_ast(i, yc_hook, {}) for i in obj.indices]
yc_hook.new_grid('dummy_grid_true', dimensions)
# Create 'hook' kernel solution
yk_hook = YaskKernel(name, yc_hook)
grid = yk_hook.new_grid(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_Space:
# 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)
else:
# time and misc dimensions
assert grid.is_dim_used(i.name)
assert grid.get_alloc_size(i.name) == s
grid.alloc_storage()
self.grids[grid.get_name()] = grid
return grid
def make_grid(self, obj):
"""
Create a Data wrapping a YASK grid. Memory is allocated.
Parameters
----------
obj : Function
The symbolic object for which a new YASK grid is created.
"""
# 'hook' compiler solution: describes the grid
# 'hook' kernel solution: allocates memory
# A unique name for the 'hook' compiler and kernel solutions
suffix = Signer._digest(self, obj, configuration, YaskContext._hookcounter)
YaskContext._hookcounter += 1
name = namespace['jit-hook'](suffix)
# Create 'hook' compiler solution
yc_hook = self.make_yc_solution(name)
if obj.indices != self.dimensions:
# Note: YASK wants *at least* a grid with *all* space (domain) dimensions
# *and* the stepping dimension. `obj`, however, may actually employ a
# different set of dimensions (e.g., a strict subset and/or some misc
# dimensions). In such a case, an extra dummy grid is attached
# `obj` examples: u(x, d), u(x, y, z)
dimensions = [make_yask_ast(i, yc_hook) for i in self.dimensions]
yc_hook.new_grid('dummy_grid_full', dimensions)
dimensions = [make_yask_ast(i.root, yc_hook) for i in obj.indices]
yc_hook.new_grid('dummy_grid_true', dimensions)
# Create 'hook' kernel solution
yk_hook = YaskKernel(name, yc_hook)
grid = yk_hook.new_grid(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._size_halo):
if i.is_Space:
# 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)
else:
# time and misc dimensions
assert grid.is_dim_used(i.root.name)
assert grid.get_alloc_size(i.root.name) == s
grid.alloc_storage()
self.grids[grid.get_name()] = grid
return grid
def _specialize_iet(self, iet, **kwargs):
"""
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.
"""
mapper = {}
self.yk_solns = OrderedDict()
for n, (section, trees) in enumerate(find_affine_trees(iet).items()):
dimensions = tuple(
filter_ordered(i.dim.root for i in flatten(trees)))
context = contexts.fetch(dimensions, self._dtype)
# A unique name for the 'real' compiler and kernel solutions
name = namespace['jit-soln'](Signer._digest(
configuration, *[i.root for i in trees]))
# Create a YASK compiler solution for this Operator
yc_soln = context.make_yc_solution(name)
try:
# Generate YASK grids and populate `yc_soln` with equations
local_grids = yaskit(trees, yc_soln)
# Build the new IET nodes
yk_soln_obj = YaskSolnObject(namespace['code-soln-name'](n))
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'], yk_soln_obj)
funcall = Offloaded(funcall, self._dtype)
mapper[trees[0].root] = funcall
mapper.update({i.root: mapper.get(i.root)
for i in trees}) # Drop trees
# Mark `funcall` as an external function call
self._func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
yk_soln = context.make_yk_solution(name, yc_soln, local_grids)
self.yk_solns[(dimensions, yk_soln_obj)] = yk_soln
# 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 NotImplementedError as e:
log("Unable to offload a candidate tree. Reason: [%s]" %
str(e))
iet = Transformer(mapper).visit(iet)
if not self.yk_solns:
log("No offloadable trees found")
# 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
yk_grid_objs = {
i.name: YaskGridObject(i.name)
for i in self._input if i.from_YASK
}
yk_grid_objs.update({i: YaskGridObject(i) for i in self._local_grids})
iet = make_grid_accesses(iet, yk_grid_objs)
# Finally optimize all non-yaskized loops
iet = super(OperatorYASK, self)._specialize_iet(iet, **kwargs)
return iet
def _soname(self):
"""A unique name for the shared object resulting from JIT compilation."""
return Signer._digest(self, configuration)
def make_var(self, obj):
"""
Create a Data wrapping a YASK var. Memory is allocated.
Parameters
----------
obj : Function
The symbolic object for which a new YASK var is created.
"""
# 'hook' compiler solution: describes the var
# 'hook' kernel solution: allocates memory
# A unique name for the 'hook' compiler and kernel solutions
suffix = Signer._digest(self, obj, configuration,
YaskContext._hookcounter)
YaskContext._hookcounter += 1
name = namespace['jit-hook'](suffix)
# Create 'hook' compiler solution
yc_hook = self.make_yc_solution(name)
# Tell YASK compiler about *all* space (domain) dimensions *and* the
# stepping dimension. This is done to ensure that all hook solutions
# have the same list of problem dimensions. Note: `obj` may
# actually employ a different set of dimensions (e.g., a strict
# subset and/or some misc dimensions).
space_dims = [make_yask_ast(i, yc_hook) for i in self.space_dimensions]
yc_hook.set_domain_dims(space_dims)
step_dim = make_yask_ast(self.step_dimension, yc_hook)
yc_hook.set_step_dim(step_dim)
# Create YASK compiler variable based on dimensions of `obj`.
# This is done to force YASK to generate code to create vars
# of this type, which will be needed when creating the YASK
# kernel variable below.
dimensions = [make_yask_ast(i.root, yc_hook) for i in obj.indices]
yc_hook.new_var('template_var', dimensions)
# Create 'hook' kernel solution from `yc_hook`
yk_hook = YaskKernel(name, yc_hook)
# Create YASK kernel variable for `obj`. YASK knows how to
# create a variable of these dimesions because of the
# 'dummy_var' declared above.
var = yk_hook.new_var(obj)
# Where should memory be allocated ?
alloc = obj._allocator
if alloc.is_Numa:
if alloc.put_onnode:
var.set_numa_preferred(alloc.node)
elif alloc.put_local:
var.set_numa_preferred(namespace['numa-put-local'])
for i, s, h in zip(obj.indices, obj.shape_allocated, obj._size_halo):
if i.is_Space:
# 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 accommodate it."
var.set_left_halo_size(i.name, h.left)
var.set_right_halo_size(i.name, h.right)
else:
# time and misc dimensions
assert var.is_dim_used(i.root.name)
assert var.get_alloc_size(i.root.name) == s
var.alloc_storage()
self.vars[var.get_name()] = var
return var
def _soname(self):
"""
A unique name for the shared object resulting from the jit-compilation
of this Operator.
"""
return Signer._digest(self, configuration)
def _specialize_iet(self, iet, **kwargs):
"""
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.
"""
offloadable = find_offloadable_trees(iet)
if len(offloadable.trees) == 0:
self.yk_soln = YaskNullKernel()
log("No offloadable trees found")
else:
context = contexts.fetch(offloadable.grid, offloadable.dtype)
# A unique name for the 'real' compiler and kernel solutions
name = namespace['jit-soln'](Signer._digest(iet, configuration))
# Create a YASK compiler solution for this Operator
yc_soln = context.make_yc_solution(name)
try:
trees = offloadable.trees
# Generate YASK grids and populate `yc_soln` with equations
mapper = yaskizer(trees, yc_soln)
local_grids = [i for i in mapper if i.is_Array]
# Transform the IET
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'],
namespace['code-soln-name'])
funcall = Element(c.Statement(ccode(funcall)))
mapper = {trees[0].root: funcall}
mapper.update({i.root: mapper.get(i.root)
for i in trees}) # Drop trees
iet = Transformer(mapper).visit(iet)
# Mark `funcall` as an external function call
self.func_table[namespace['code-soln-run']] = MetaCall(
None, False)
# JIT-compile the newly-created YASK kernel
self.yk_soln = context.make_yk_solution(
name, 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 NotImplementedError as e:
self.yk_soln = YaskNullKernel()
log("Unable to offload a candidate tree. Reason: [%s]" %
str(e))
# 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
iet = make_grid_accesses(iet)
# Finally optimize all non-yaskized loops
iet = super(Operator, self)._specialize_iet(iet, **kwargs)
return iet
def make_yask_kernels(iet, **kwargs):
yk_solns = kwargs.pop('yk_solns')
mapper = {}
for n, (section, trees) in enumerate(find_affine_trees(iet).items()):
dimensions = tuple(filter_ordered(i.dim.root for i in flatten(trees)))
# Retrieve the section dtype
exprs = FindNodes(Expression).visit(section)
dtypes = {e.dtype for e in exprs}
if len(dtypes) != 1:
log("Unable to offload in presence of mixed-precision arithmetic")
continue
dtype = dtypes.pop()
context = contexts.fetch(dimensions, dtype)
# A unique name for the 'real' compiler and kernel solutions
name = namespace['jit-soln'](Signer._digest(configuration,
*[i.root for i in trees]))
# Create a YASK compiler solution for this Operator
yc_soln = context.make_yc_solution(name)
try:
# Generate YASK vars and populate `yc_soln` with equations
local_vars = yaskit(trees, yc_soln)
# Build the new IET nodes
yk_soln_obj = YASKSolnObject(namespace['code-soln-name'](n))
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'], yk_soln_obj)
funcall = Offloaded(funcall, dtype)
mapper[trees[0].root] = funcall
mapper.update({i.root: mapper.get(i.root)
for i in trees}) # Drop trees
# JIT-compile the newly-created YASK kernel
yk_soln = context.make_yk_solution(name, yc_soln, local_vars)
yk_solns[(dimensions, yk_soln_obj)] = yk_soln
# Print some useful information about the newly constructed solution
log("Solution '%s' contains %d var(s) and %d equation(s)." %
(yc_soln.get_name(), yc_soln.get_num_vars(),
yc_soln.get_num_equations()))
except NotImplementedError as e:
log("Unable to offload a candidate tree. Reason: [%s]" % str(e))
iet = Transformer(mapper).visit(iet)
if not yk_solns:
log("No offloadable trees found")
# Some Iteration/Expression trees are not offloaded to YASK and may
# require further processing to be executed through YASK, due to the
# different storage layout
yk_var_objs = {
i.name: YASKVarObject(i.name)
for i in FindSymbols().visit(iet) if i.from_YASK
}
yk_var_objs.update({i: YASKVarObject(i) for i in get_local_vars(yk_solns)})
iet = make_var_accesses(iet, yk_var_objs)
# The signature needs to be updated
# TODO: this could be done automagically through the iet pass engine, but
# currently it only supports *appending* to the parameters list. While here
# we actually need to change it as some parameters may disappear (x_m, x_M, ...)
parameters = derive_parameters(iet, True)
iet = iet._rebuild(parameters=parameters)
return iet, {}
def _specialize_iet(self, iet, **kwargs):
"""
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.
"""
mapper = {}
self.yk_solns = OrderedDict()
for n, (section, trees) in enumerate(find_affine_trees(iet).items()):
dimensions = tuple(filter_ordered(i.dim.root for i in flatten(trees)))
context = contexts.fetch(dimensions, self._dtype)
# A unique name for the 'real' compiler and kernel solutions
name = namespace['jit-soln'](Signer._digest(configuration,
*[i.root for i in trees]))
# Create a YASK compiler solution for this Operator
yc_soln = context.make_yc_solution(name)
try:
# Generate YASK grids and populate `yc_soln` with equations
local_grids = yaskit(trees, yc_soln)
# Build the new IET nodes
yk_soln_obj = YaskSolnObject(namespace['code-soln-name'](n))
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'], yk_soln_obj)
funcall = Offloaded(funcall, self._dtype)
mapper[trees[0].root] = funcall
mapper.update({i.root: mapper.get(i.root) for i in trees}) # Drop trees
# Mark `funcall` as an external function call
self._func_table[namespace['code-soln-run']] = MetaCall(None, False)
# JIT-compile the newly-created YASK kernel
yk_soln = context.make_yk_solution(name, yc_soln, local_grids)
self.yk_solns[(dimensions, yk_soln_obj)] = yk_soln
# 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 NotImplementedError as e:
log("Unable to offload a candidate tree. Reason: [%s]" % str(e))
iet = Transformer(mapper).visit(iet)
if not self.yk_solns:
log("No offloadable trees found")
# 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
yk_grid_objs = {i.name: YaskGridObject(i.name) for i in self._input
if i.from_YASK}
yk_grid_objs.update({i: YaskGridObject(i) for i in self._local_grids})
iet = make_grid_accesses(iet, yk_grid_objs)
# Finally optimize all non-yaskized loops
iet = super(OperatorYASK, self)._specialize_iet(iet, **kwargs)
return iet
def _soname(self):
"""A unique name for the shared object resulting from JIT compilation."""
return Signer._digest(self, configuration)
def _specialize_iet(cls, iet, **kwargs):
"""
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.
"""
mapper = {}
yk_solns = kwargs.pop('yk_solns')
for n, (section, trees) in enumerate(find_affine_trees(iet).items()):
dimensions = tuple(
filter_ordered(i.dim.root for i in flatten(trees)))
# Retrieve the section dtype
exprs = FindNodes(Expression).visit(section)
dtypes = {e.dtype for e in exprs}
if len(dtypes) != 1:
log("Unable to offload in presence of mixed-precision arithmetic"
)
continue
dtype = dtypes.pop()
context = contexts.fetch(dimensions, dtype)
# A unique name for the 'real' compiler and kernel solutions
name = namespace['jit-soln'](Signer._digest(
configuration, *[i.root for i in trees]))
# Create a YASK compiler solution for this Operator
yc_soln = context.make_yc_solution(name)
try:
# Generate YASK vars and populate `yc_soln` with equations
local_vars = yaskit(trees, yc_soln)
# Build the new IET nodes
yk_soln_obj = YaskSolnObject(namespace['code-soln-name'](n))
funcall = make_sharedptr_funcall(namespace['code-soln-run'],
['time'], yk_soln_obj)
funcall = Offloaded(funcall, dtype)
mapper[trees[0].root] = funcall
mapper.update({i.root: mapper.get(i.root)
for i in trees}) # Drop trees
# JIT-compile the newly-created YASK kernel
yk_soln = context.make_yk_solution(name, yc_soln, local_vars)
yk_solns[(dimensions, yk_soln_obj)] = yk_soln
# Print some useful information about the newly constructed solution
log("Solution '%s' contains %d var(s) and %d equation(s)." %
(yc_soln.get_name(), yc_soln.get_num_vars(),
yc_soln.get_num_equations()))
except NotImplementedError as e:
log("Unable to offload a candidate tree. Reason: [%s]" %
str(e))
iet = Transformer(mapper).visit(iet)
if not yk_solns:
log("No offloadable trees found")
# Some Iteration/Expression trees are not offloaded to YASK and may
# require further processing to be executed through YASK, due to the
# different storage layout
yk_var_objs = {
i.name: YaskVarObject(i.name)
for i in FindSymbols().visit(iet) if i.from_YASK
}
yk_var_objs.update(
{i: YaskVarObject(i)
for i in cls._get_local_vars(yk_solns)})
iet = make_var_accesses(iet, yk_var_objs)
# The signature needs to be updated
parameters = derive_parameters(iet, True)
iet = iet._rebuild(parameters=parameters)
return super(OperatorYASK, cls)._specialize_iet(iet, **kwargs)