def dimension_sort(expr, key=None):
"""
Topologically sort the :class:`Dimension`s in ``expr``, based on the order
in which they are encountered when visiting ``expr``.
:param expr: The :class:`sympy.Eq` from which the :class:`Dimension`s are
extracted. They can appear both as array indices or as free
symbols.
:param key: A callable used as key to enforce a final ordering.
"""
# Get the Indexed dimensions, in appearance order
constraints = [
tuple(i.indices) for i in retrieve_indexed(expr, mode='all')
]
for i, constraint in enumerate(list(constraints)):
normalized = []
for j in constraint:
found = [d for d in j.free_symbols if isinstance(d, Dimension)]
normalized.extend([d for d in found if d not in normalized])
constraints[i] = normalized
ordering = partial_order(constraints)
# Add any leftover free dimensions (not an Indexed' index)
dimensions = [i for i in expr.free_symbols if isinstance(i, Dimension)]
dimensions = filter_sorted(dimensions,
key=attrgetter('name')) # for determinism
ordering.extend([i for i in dimensions if i not in ordering])
# Add parent dimensions
derived = [i for i in ordering if i.is_Derived]
for i in derived:
ordering.insert(ordering.index(i), i.parent)
return sorted(ordering, key=lambda i: not i.is_Time)
def xreplace_indices(exprs, mapper, key=None, only_rhs=False):
"""
Replace array indices in expressions.
Parameters
----------
exprs : expr-like or list of expr-like
One or more expressions to which the replacement is applied.
mapper : dict
The substitution rules.
key : list of symbols or callable
An escape hatch to rule out some objects from the replacement.
If a list, apply the replacement to the symbols in ``key`` only. If a
callable, apply the replacement to a symbol S if and only if ``key(S)``
gives True.
only_rhs : bool, optional
If True, apply the replacement to Eq right-hand sides only.
"""
get = lambda i: i.rhs if only_rhs is True else i
handle = flatten(retrieve_indexed(get(i)) for i in as_tuple(exprs))
if isinstance(key, Iterable):
handle = [i for i in handle if i.base.label in key]
elif callable(key):
handle = [i for i in handle if key(i)]
mapper = dict(zip(handle, [i.xreplace(mapper) for i in handle]))
replaced = [uxreplace(i, mapper) for i in as_tuple(exprs)]
return replaced if isinstance(exprs, Iterable) else replaced[0]
def dimension_sort(expr, key=None):
"""
Topologically sort the :class:`Dimension`s in ``expr``, based on the order
in which they are encountered when visiting ``expr``.
:param expr: The :class:`sympy.Eq` from which the :class:`Dimension`s are
extracted. They can appear both as array indices or as free
symbols.
:param key: A callable used as key to enforce a final ordering.
"""
# Get all Indexed dimensions, in the same order as the appear in /expr/
constraints = []
for i in retrieve_indexed(expr, mode='all'):
constraint = []
for ai, fi in zip(i.indices, i.base.function.indices):
if ai.is_Number:
constraint.append(fi)
else:
constraint.extend([d for d in ai.free_symbols
if isinstance(d, Dimension) and d not in constraint])
constraints.append(tuple(constraint))
ordering = partial_order(constraints)
# Add any leftover free dimensions (not an Indexed' index)
dimensions = [i for i in expr.free_symbols if isinstance(i, Dimension)]
dimensions = filter_sorted(dimensions, key=attrgetter('name')) # for determinism
ordering.extend([i for i in dimensions if i not in ordering])
# Add parent dimensions
derived = [i for i in ordering if i.is_Derived]
for i in derived:
ordering.insert(ordering.index(i), i.parent)
return sorted(ordering, key=lambda i: not i.is_Time)
def xreplace_indices(exprs, mapper, key=None, only_rhs=False):
"""
Replace array indices in expressions.
Parameters
----------
exprs : expr-like or list of expr-like
One or more expressions to which the replacement is applied.
mapper : dict
The substitution rules.
key : list of symbols or callable
An escape hatch to rule out some objects from the replacement.
If a list, apply the replacement to the symbols in ``key`` only. If a
callable, apply the replacement to a symbol S if and only if ``key(S)``
gives True.
only_rhs : bool, optional
If True, apply the replacement to Eq right-hand sides only.
"""
get = lambda i: i.rhs if only_rhs is True else i
handle = flatten(retrieve_indexed(get(i)) for i in as_tuple(exprs))
if isinstance(key, Iterable):
handle = [i for i in handle if i.base.label in key]
elif callable(key):
handle = [i for i in handle if key(i)]
mapper = dict(zip(handle, [i.xreplace(mapper) for i in handle]))
replaced = [i.xreplace(mapper) for i in as_tuple(exprs)]
return replaced if isinstance(exprs, Iterable) else replaced[0]
def estimate_memory(handle, mode='realistic'):
"""
Estimate the number of memory reads and writes.
:param handle: a SymPy expression or an iterator of SymPy expressions.
:param mode: Mode for computing the estimate:
Estimate ``mode`` might be any of: ::
* ideal: Also known as "compulsory traffic", which is the minimum
number of read/writes to be performed (ie, models an infinite cache).
* ideal_with_stores: Like ideal, but a data item which is both read.
and written is counted twice (ie both load an
store are counted).
* realistic: Assume that all datasets, even the time-independent ones,
need to be re-read at each time iteration.
"""
assert mode in ['ideal', 'ideal_with_stores', 'realistic']
def access(symbol):
assert isinstance(symbol, Indexed)
# Irregular accesses (eg A[B[i]]) are counted as compulsory traffic
if any(i.atoms(Indexed) for i in symbol.indices):
return symbol
else:
return symbol.base
try:
# Is it a plain SymPy object ?
iter(handle)
except TypeError:
handle = [handle]
if mode in ['ideal', 'ideal_with_stores']:
filter = lambda s: any(q_timedimension(i) for i in s.atoms())
else:
filter = lambda s: s
reads = set(flatten([retrieve_indexed(e.rhs) for e in handle]))
writes = set(flatten([retrieve_indexed(e.lhs) for e in handle]))
reads = set([access(s) for s in reads if filter(s)])
writes = set([access(s) for s in writes if filter(s)])
if mode == 'ideal':
return len(set(reads) | set(writes))
else:
return len(reads) + len(writes)
def q_indirect(expr):
"""
Return True if ``indexed`` has indirect accesses, False otherwise.
:Examples:
a[i] --> False
a[b[i]] --> True
"""
from devito.symbolics.search import retrieve_indexed
if not expr.is_Indexed:
return False
return any(retrieve_indexed(i) for i in expr.indices)
def xreplace_indices(exprs, mapper, candidates=None, only_rhs=False):
"""
Create new expressions from ``exprs``, by replacing all index variables
specified in mapper appearing as a tensor index. Only tensors whose symbolic
name appears in ``candidates`` are considered if ``candidates`` is not None.
"""
get = lambda i: i.rhs if only_rhs is True else i
handle = flatten(retrieve_indexed(get(i)) for i in as_tuple(exprs))
if candidates is not None:
handle = [i for i in handle if i.base.label in candidates]
mapper = dict(zip(handle, [i.xreplace(mapper) for i in handle]))
replaced = [i.xreplace(mapper) for i in as_tuple(exprs)]
return replaced if isinstance(exprs, Iterable) else replaced[0]
def q_indirect(expr):
"""
Return True if ``indexed`` has indirect accesses, False otherwise.
:Examples:
a[i] --> False
a[b[i]] --> True
"""
from devito.symbolics.search import retrieve_indexed
if not expr.is_Indexed:
return False
return any(retrieve_indexed(i) for i in expr.indices)
def xreplace_indices(exprs, mapper, key=None, only_rhs=False):
"""
Replace indices in SymPy equations.
:param exprs: The target SymPy expression, or a collection of SymPy expressions.
:param mapper: A dictionary containing the index substitution rules.
:param key: (Optional) either an iterable or a function. In the former case,
all objects whose name does not appear in ``key`` are ruled out.
Likewise, if a function, all objects for which ``key(obj)`` gives
False are ruled out.
:param only_rhs: (Optional) apply the substitution rules to right-hand sides only.
"""
get = lambda i: i.rhs if only_rhs is True else i
handle = flatten(retrieve_indexed(get(i)) for i in as_tuple(exprs))
if isinstance(key, Iterable):
handle = [i for i in handle if i.base.label in key]
elif callable(key):
handle = [i for i in handle if key(i)]
mapper = dict(zip(handle, [i.xreplace(mapper) for i in handle]))
replaced = [i.xreplace(mapper) for i in as_tuple(exprs)]
return replaced if isinstance(exprs, Iterable) else replaced[0]
