def topological_sort(exprs):
"""Topologically sort the temporaries in a list of equations."""
mapper = {e.lhs: e for e in exprs}
assert len(mapper) == len(exprs) # Expect SSA
# Build DAG and topologically-sort temporaries
temporaries, tensors = split(exprs, lambda e: not e.lhs.is_Indexed)
dag = DAG(nodes=temporaries)
for e in temporaries:
for r in retrieve_terminals(e.rhs):
if r not in mapper:
continue
elif mapper[r] is e:
# Avoid cyclic dependences, such as
# Eq(f, f + 1)
continue
elif r.is_Indexed:
# Only scalars enforce an ordering
continue
else:
dag.add_edge(mapper[r], e, force_add=True)
processed = dag.topological_sort()
# Append tensor equations at the end in user-provided order
processed.extend(tensors)
return processed
def process(func, state):
"""
Apply ``func`` to the IETs in ``state._efuncs``, and update ``state`` accordingly.
"""
# Create a Call graph. `func` will be applied to each node in the Call graph.
# `func` might change an `efunc` signature; the Call graph will be used to
# propagate such change through the `efunc` callers
dag = DAG(nodes=['root'])
queue = ['root']
while queue:
caller = queue.pop(0)
callees = FindNodes(Call).visit(state._efuncs[caller])
for callee in filter_ordered([i.name for i in callees]):
if callee in state._efuncs: # Exclude foreign Calls, e.g., MPI calls
try:
dag.add_node(callee)
queue.append(callee)
except KeyError:
# `callee` already in `dag`
pass
dag.add_edge(callee, caller)
assert dag.size == len(state._efuncs)
# Apply `func`
for i in dag.topological_sort():
state._efuncs[i], metadata = func(state._efuncs[i])
# Track any new Dimensions introduced by `func`
state._dimensions.extend(list(metadata.get('dimensions', [])))
# Track any new #include required by `func`
state._includes.extend(list(metadata.get('includes', [])))
state._includes = filter_ordered(state._includes)
# Track any new ElementalFunctions
state._efuncs.update(
OrderedDict([(i.name, i) for i in metadata.get('efuncs', [])]))
# If there's a change to the `args` and the `iet` is an efunc, then
# we must update the call sites as well, as the arguments dropped down
# to the efunc have just increased
args = as_tuple(metadata.get('args'))
if args:
# `extif` avoids redundant updates to the parameters list, due
# to multiple children wanting to add the same input argument
extif = lambda v: list(v) + [e for e in args if e not in v]
stack = [i] + dag.all_downstreams(i)
for n in stack:
efunc = state._efuncs[n]
calls = [
c for c in FindNodes(Call).visit(efunc) if c.name in stack
]
mapper = {
c: c._rebuild(arguments=extif(c.arguments))
for c in calls
}
efunc = Transformer(mapper).visit(efunc)
if efunc.is_Callable:
efunc = efunc._rebuild(parameters=extif(efunc.parameters))
state._efuncs[n] = efunc
def process(func, state):
"""
Apply ``func`` to the IETs in ``state._efuncs``, and update ``state`` accordingly.
"""
# Create a Call graph. `func` will be applied to each node in the Call graph.
# `func` might change an `efunc` signature; the Call graph will be used to
# propagate such change through the `efunc` callers
dag = DAG(nodes=['root'])
queue = ['root']
while queue:
caller = queue.pop(0)
callees = FindNodes(Call).visit(state._efuncs[caller])
for callee in filter_ordered([i.name for i in callees]):
if callee in state._efuncs: # Exclude foreign Calls, e.g., MPI calls
try:
dag.add_node(callee)
queue.append(callee)
except KeyError:
# `callee` already in `dag`
pass
dag.add_edge(callee, caller)
assert dag.size == len(state._efuncs)
# Apply `func`
for i in dag.topological_sort():
state._efuncs[i], metadata = func(state._efuncs[i])
# Track any new Dimensions introduced by `func`
state._dimensions.extend(list(metadata.get('dimensions', [])))
# Track any new #include required by `func`
state._includes.extend(list(metadata.get('includes', [])))
state._includes = filter_ordered(state._includes)
# Track any new ElementalFunctions
state._efuncs.update(OrderedDict([(i.name, i)
for i in metadata.get('efuncs', [])]))
# If there's a change to the `args` and the `iet` is an efunc, then
# we must update the call sites as well, as the arguments dropped down
# to the efunc have just increased
args = as_tuple(metadata.get('args'))
if args:
# `extif` avoids redundant updates to the parameters list, due
# to multiple children wanting to add the same input argument
extif = lambda v: list(v) + [e for e in args if e not in v]
stack = [i] + dag.all_downstreams(i)
for n in stack:
efunc = state._efuncs[n]
calls = [c for c in FindNodes(Call).visit(efunc) if c.name in stack]
mapper = {c: c._rebuild(arguments=extif(c.arguments)) for c in calls}
efunc = Transformer(mapper).visit(efunc)
if efunc.is_Callable:
efunc = efunc._rebuild(parameters=extif(efunc.parameters))
state._efuncs[n] = efunc
def topological_sort(exprs):
"""Topologically sort a list of equations."""
mapper = {e.lhs: e for e in exprs}
assert len(mapper) == len(exprs) # Expect SSA
dag = DAG(nodes=exprs)
for e in exprs:
for r in retrieve_terminals(e.rhs):
if r not in mapper:
continue
elif mapper[r] is e:
# Avoid cyclic dependences, such as
# Eq(f, f + 1)
continue
elif r.is_Indexed:
# Only scalars enforce an ordering
continue
else:
dag.add_edge(mapper[r], e, force_add=True)
processed = dag.topological_sort()
return processed
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
grids = {getattr(kwargs[p.name], 'grid', None) for p in overrides}
grids.update({getattr(p, 'grid', None) for p in defaults})
grids.discard(None)
if len(grids) > 1 and configuration['mpi']:
raise ValueError("Multiple Grids found")
try:
grid = grids.pop()
args.update(grid._arg_values(**kwargs))
except KeyError:
grid = None
# Process Dimensions
# A topological sorting is used so that derived Dimensions are processed after
# their parents (note that a leaf Dimension can have an arbitrary long list of
# ancestors)
dag = DAG(self.dimensions,
[(i, i.parent) for i in self.dimensions if i.is_Derived])
for d in reversed(dag.topological_sort()):
args.update(d._arg_values(args, self._dspace[d], grid, **kwargs))
# Process Objects (which may need some `args`)
for o in self.objects:
args.update(o._arg_values(args, grid=grid, **kwargs))
# Sanity check
for p in self.parameters:
p._arg_check(args, self._dspace[p])
for d in self.dimensions:
if d.is_Derived:
d._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.parameters:
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 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))
# Attach `grid` to the arguments map
args = ArgumentsMap(grid, **args)
return args
def _prepare_arguments(self, autotune=None, **kwargs):
"""
Process runtime arguments passed to ``.apply()` and derive
default values for any remaining arguments.
"""
# Sanity check -- all user-provided keywords must be 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))
overrides, defaults = split(self.input, lambda p: p.name in kwargs)
# Process data-carrier overrides
args = kwargs['args'] = ReducerMap()
for p in overrides:
args.update(p._arg_values(**kwargs))
try:
args.reduce_inplace()
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 = kwargs['args'] = args.reduce_all()
# DiscreteFunctions may be created from CartesianDiscretizations, which in
# turn could be Grids or SubDomains. Both may provide arguments
discretizations = {
getattr(kwargs[p.name], 'grid', None)
for p in overrides
}
discretizations.update({getattr(p, 'grid', None) for p in defaults})
discretizations.discard(None)
for i in discretizations:
args.update(i._arg_values(**kwargs))
# There can only be one Grid from which DiscreteFunctions were created
grids = {i for i in discretizations if isinstance(i, Grid)}
if len(grids) > 1:
# We loosely tolerate multiple Grids for backwards compatibility
# with spacial subsampling, which should be revisited however. And
# With MPI it would definitely break!
if configuration['mpi']:
raise ValueError("Multiple Grids found")
try:
grid = grids.pop()
except KeyError:
grid = None
# An ArgumentsMap carries additional metadata that may be used by
# the subsequent phases of the arguments processing
args = kwargs['args'] = ArgumentsMap(args, grid, self._allocator,
self._platform)
# Process Dimensions
# A topological sorting is used so that derived Dimensions are processed after
# their parents (note that a leaf Dimension can have an arbitrary long list of
# ancestors)
dag = DAG(self.dimensions,
[(i, i.parent) for i in self.dimensions if i.is_Derived])
for d in reversed(dag.topological_sort()):
args.update(d._arg_values(self._dspace[d], grid, **kwargs))
# Process Objects
for o in self.objects:
args.update(o._arg_values(grid=grid, **kwargs))
# In some "lower-level" Operators implementing a random piece of C, such as
# one or more calls to third-party library functions, there could still be
# at this point unprocessed arguments (e.g., scalars)
kwargs.pop('args')
args.update({k: v for k, v in kwargs.items() if k not in args})
# Sanity check
for p in self.parameters:
p._arg_check(args, self._dspace[p])
for d in self.dimensions:
if d.is_Derived:
d._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.parameters:
try:
args.update(kwargs.get(p.name, p)._arg_finalize(args, alias=p))
except AttributeError:
# User-provided floats/ndarray obviously do not have `_arg_finalize`
args.update(p._arg_finalize(args, alias=p))
# Execute autotuning and adjust arguments accordingly
args.update(
self._autotune(args, autotune or configuration['autotuning']))
return args
class State(object):
def __init__(self, iet):
self._efuncs = OrderedDict([('main', iet)])
self._dimensions = []
self._input = []
self._includes = []
self._call_graph = DAG(nodes=['main'])
def _process(self, func):
"""Apply ``func`` to all tracked ``IETs``."""
for i in self._call_graph.topological_sort():
self._efuncs[i], metadata = func(self._efuncs[i])
# Track any new Dimensions and includes introduced by `func`
self._dimensions.extend(list(metadata.get('dimensions', [])))
self._includes.extend(list(metadata.get('includes', [])))
# If there's a change to the `input` and the `iet` is an efunc, then
# we must update the call sites as well, as the arguments dropped down
# to the efunc have just increased
_input = as_tuple(metadata.get('input'))
if _input:
# `extif` avoids redundant updates to the parameters list, due
# to multiple children wanting to add the same input argument
extif = lambda v: list(v) + [e for e in _input if e not in v]
stack = [i] + self._call_graph.all_downstreams(i)
for n in stack:
efunc = self._efuncs[n]
calls = [c for c in FindNodes(Call).visit(efunc) if c.name in stack]
mapper = {c: c._rebuild(arguments=extif(c.arguments)) for c in calls}
efunc = Transformer(mapper).visit(efunc)
if efunc.is_Callable:
efunc = efunc._rebuild(parameters=extif(efunc.parameters))
self._efuncs[n] = efunc
self._input.extend(list(_input))
for k, v in metadata.get('efuncs', {}).items():
# Update the efuncs
if k.is_Callable:
self._efuncs[k.name] = k
# Update the call graph
self._call_graph.add_node(k.name, ignore_existing=True)
for target in (v or [None]):
self._call_graph.add_edge(k.name, target or 'main', force_add=True)
@property
def root(self):
return self._efuncs['main']
@property
def efuncs(self):
return tuple(v for k, v in self._efuncs.items() if k != 'main')
@property
def dimensions(self):
return self._dimensions
@property
def input(self):
return self._input
@property
def includes(self):
return self._includes