def find_offloadable_trees(iet):
"""
Return the trees within ``iet`` 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 = namedtuple('Offlodable', 'trees grid dtype')
Offloadable.__new__.__defaults__ = [], None, None
reducer = ReducerMap()
# Find offloadable candidates
for tree in retrieve_iteration_tree(iet):
# The outermost iteration must be over time and it must
# nest at least one iteration
if len(tree) <= 1:
continue
if not tree.root.dim.is_Time:
continue
grid_tree = tree[1:]
if not all(i.is_Affine for i in grid_tree):
# Non-affine array accesses unsupported by YASK
continue
bundles = [i for i in tree.inner.nodes if i.is_ExpressionBundle]
if len(bundles) != 1:
# Illegal nest
continue
bundle = bundles[0]
# Found an offloadable candidate
reducer.setdefault('grid_trees', []).append(grid_tree)
# Track `grid` and `dtype`
functions = flatten(i.functions for i in bundle.exprs)
reducer.extend(
('grid', i.grid) for i in functions if i.is_TimeFunction)
reducer.extend(
('dtype', i.dtype) for i in functions if i.is_TimeFunction)
# `grid` and `dtype` must be unique
try:
grid = reducer.unique('grid')
dtype = reducer.unique('dtype')
trees = reducer['grid_trees']
except (KeyError, ValueError):
return Offloadable()
# Do the trees iterate over a convex section of the grid?
if not all(i.dim._defines | set(grid.dimensions) for i in flatten(trees)):
return Offloadable()
return Offloadable(trees, grid, dtype)
def _prepare_arguments(self, **kwargs):
"""
Process runtime arguments passed to ``.apply()` and derive
default values for any remaining arguments.
"""
# Process data-carriers (first overrides, then fill up with whatever is needed)
args = ReducerMap()
args.update(
[p._arg_values(**kwargs) for p in self.input if p.name in kwargs])
args.update(
[p._arg_values() for p in self.input if p.name not in args])
args = args.reduce_all()
# All TensorFunctions should be defined on the same Grid
functions = [
kwargs.get(p, p) for p in self.input if p.is_TensorFunction
]
mapper = ReducerMap([('grid', i.grid) for i in functions if i.grid])
try:
grid = mapper.unique('grid')
except (KeyError, ValueError):
if mapper and configuration['mpi']:
raise RuntimeError("Multiple `Grid`s found before `apply`")
grid = None
# Process dimensions (derived go after as they might need/affect their parents)
derived, main = split(self.dimensions, lambda i: i.is_Derived)
for p in main:
args.update(p._arg_values(args, self._dspace[p], grid, **kwargs))
for p in derived:
args.update(p._arg_values(args, self._dspace[p], grid, **kwargs))
# Sanity check
for p in self.input:
p._arg_check(args, self._dspace[p])
# Add in the profiler argument
args[self._profiler.name] = self._profiler.timer.reset()
# Add in any backend-specific argument
args.update(kwargs.pop('backend', {}))
# Execute autotuning and adjust arguments accordingly
args = self._autotune(
args, kwargs.pop('autotune', configuration['autotuning']))
# Check all user-provided keywords are known to the Operator
if not configuration['ignore-unknowns']:
for k, v in kwargs.items():
if k not in self._known_arguments:
raise ValueError("Unrecognized argument %s=%s" % (k, v))
return args
def _prepare_arguments(self, **kwargs):
"""
Process runtime arguments passed to ``.apply()` and derive
default values for any remaining arguments.
"""
# Process data-carriers (first overrides, then fill up with whatever is needed)
args = ReducerMap()
args.update(
[p._arg_values(**kwargs) for p in self.input if p.name in kwargs])
args.update(
[p._arg_values() for p in self.input if p.name not in args])
args = args.reduce_all()
# All TensorFunctions should be defined on the same Grid
functions = [
kwargs.get(p, p) for p in self.input if p.is_TensorFunction
]
mapper = ReducerMap([('grid', i.grid) for i in functions if i.grid])
try:
grid = mapper.unique('grid')
except (KeyError, ValueError):
if mapper and configuration['mpi']:
raise RuntimeError("Multiple `Grid`s found before `apply`")
grid = None
# Process dimensions (derived go after as they might need/affect their parents)
derived, main = split(self.dimensions, lambda i: i.is_Derived)
for p in main:
args.update(p._arg_values(args, self._dspace[p], grid, **kwargs))
for p in derived:
args.update(p._arg_values(args, self._dspace[p], grid, **kwargs))
# Sanity check
for p in self.input:
p._arg_check(args, self._dspace[p])
# Derive additional values for DLE arguments
# TODO: This is not pretty, but it works for now. Ideally, the
# DLE arguments would be massaged into the IET so as to comply
# with the rest of the argument derivation procedure.
for arg in self._dle_args:
dim = arg.argument
osize = (1 + arg.original_dim.symbolic_end -
arg.original_dim.symbolic_start).subs(args)
if arg.value is None:
args[dim.symbolic_size.name] = osize
elif isinstance(arg.value, int):
args[dim.symbolic_size.name] = arg.value
else:
args[dim.symbolic_size.name] = arg.value(osize)
# Add in the profiler argument
args[self.profiler.name] = self.profiler.timer.reset()
# Add in any backend-specific argument
args.update(kwargs.pop('backend', {}))
# Execute autotuning and adjust arguments accordingly
if kwargs.pop('autotune', configuration['autotuning'].level):
args = self._autotune(args)
# Check all user-provided keywords are known to the Operator
for k, v in kwargs.items():
if k not in self._known_arguments:
raise ValueError(
"Unrecognized argument %s=%s passed to `apply`" % (k, v))
return args
def _prepare_arguments(self, **kwargs):
"""
Process runtime arguments passed to ``.apply()` and derive
default values for any remaining arguments.
"""
overrides, defaults = split(self.input, lambda p: p.name in kwargs)
# Process data-carrier overrides
args = ReducerMap()
for p in overrides:
args.update(p._arg_values(**kwargs))
try:
args = ReducerMap(args.reduce_all())
except ValueError:
raise ValueError(
"Override `%s` is incompatible with overrides `%s`" %
(p, [i for i in overrides if i.name in args]))
# Process data-carrier defaults
for p in defaults:
if p.name in args:
# E.g., SubFunctions
continue
for k, v in p._arg_values(**kwargs).items():
if k in args and args[k] != v:
raise ValueError(
"Default `%s` is incompatible with other args as "
"`%s=%s`, while `%s=%s` is expected. Perhaps you "
"forgot to override `%s`?" % (p, k, v, k, args[k], p))
args[k] = v
args = args.reduce_all()
# All DiscreteFunctions should be defined on the same Grid
functions = [
kwargs.get(p, p) for p in self.input if p.is_DiscreteFunction
]
mapper = ReducerMap([('grid', i.grid) for i in functions if i.grid])
try:
grid = mapper.unique('grid')
except (KeyError, ValueError):
if mapper and configuration['mpi']:
raise RuntimeError("Multiple `Grid`s found before `apply`")
grid = None
# Process dimensions (derived go after as they might need/affect their parents)
derived, main = split(self.dimensions, lambda i: i.is_Derived)
for p in main:
args.update(p._arg_values(args, self._dspace[p], grid, **kwargs))
for p in derived:
args.update(p._arg_values(args, self._dspace[p], grid, **kwargs))
# Sanity check
for p in self.input:
p._arg_check(args, self._dspace[p])
# Turn arguments into a format suitable for the generated code
# E.g., instead of NumPy arrays for Functions, the generated code expects
# pointers to ctypes.Struct
for p in self.input:
try:
args.update(kwargs.get(p.name, p)._arg_as_ctype(args, alias=p))
except AttributeError:
# User-provided floats/ndarray obviously do not have `_arg_as_ctype`
args.update(p._arg_as_ctype(args, alias=p))
# Add in the profiler argument
args[self._profiler.name] = self._profiler.timer.reset()
# Add in any backend-specific argument
args.update(kwargs.pop('backend', {}))
# Execute autotuning and adjust arguments accordingly
args = self._autotune(
args, kwargs.pop('autotune', configuration['autotuning']))
# Check all user-provided keywords are known to the Operator
if not configuration['ignore-unknowns']:
for k, v in kwargs.items():
if k not in self._known_arguments:
raise ValueError("Unrecognized argument %s=%s" % (k, v))
return args
