def iet_lower_dimensions(iet):
"""
Replace all DerivedDimensions within the ``iet``'s expressions with
lower-level symbolic objects (other Dimensions or Symbols).
* Array indices involving SteppingDimensions are turned into ModuloDimensions.
Example: ``u[t+1, x] = u[t, x] + 1 >>> u[t1, x] = u[t0, x] + 1``
* Array indices involving ConditionalDimensions used are turned into
integer-division expressions.
Example: ``u[t_sub, x] = u[time, x] >>> u[time / 4, x] = u[time, x]``
"""
# Lower SteppingDimensions
for i in FindNodes(Iteration).visit(iet):
if not i.uindices:
# Be quick: avoid uselessy reconstructing nodes
continue
# In an expression, there could be `u[t+1, ...]` and `v[t+1, ...]`, where
# `u` and `v` are TimeFunction with circular time buffers (save=None) *but*
# different modulo extent. The `t+1` indices above are therefore conceptually
# different, so they will be replaced with the proper ModuloDimension through
# two different calls to `xreplace`
groups = as_mapper(i.uindices, lambda d: d.modulo)
for k, v in groups.items():
mapper = {d.origin: d for d in v}
rule = lambda i: i.function.is_TimeFunction and i.function._time_size == k
replacer = lambda i: xreplace_indices(i, mapper, rule)
iet = XSubs(replacer=replacer).visit(iet)
# Lower ConditionalDimensions
cdims = [d for d in FindSymbols('free-symbols').visit(iet)
if isinstance(d, ConditionalDimension)]
mapper = {d: IntDiv(d.index, d.factor) for d in cdims}
iet = XSubs(mapper).visit(iet)
return iet
def guard(clusters):
"""
Return a new :class:`ClusterGroup` including new :class:`PartialCluster`s
for each conditional expression encountered in ``clusters``.
"""
processed = ClusterGroup()
for c in clusters:
# Find out what expressions in /c/ should be guarded
mapper = {}
for e in c.exprs:
for k, v in e.ispace.sub_iterators.items():
for i in v:
if i.dim.is_Conditional:
mapper.setdefault(i.dim, []).append(e)
# Build conditional expressions to guard clusters
conditions = {d: CondEq(d.parent % d.factor, 0) for d in mapper}
negated = {d: CondNe(d.parent % d.factor, 0) for d in mapper}
# Expand with guarded clusters
combs = list(powerset(mapper))
for dims, ndims in zip(combs, reversed(combs)):
banned = flatten(v for k, v in mapper.items() if k not in dims)
exprs = [
e.xreplace({i: IntDiv(i.parent, i.factor)
for i in mapper}) for e in c.exprs
if e not in banned
]
guards = [(i.parent, conditions[i]) for i in dims]
guards.extend([(i.parent, negated[i]) for i in ndims])
cluster = PartialCluster(exprs, c.ispace, c.dspace, c.atomics,
dict(guards))
processed.append(cluster)
return processed
def test_issue_1592(self):
grid = Grid(shape=(11, 11))
time = grid.time_dim
time_sub = ConditionalDimension('t_sub', parent=time, factor=2)
v = TimeFunction(name="v", grid=grid, space_order=4, time_dim=time_sub, save=5)
w = Function(name="w", grid=grid, space_order=4)
Operator(Eq(w, v.dx))(time=6)
op = Operator(Eq(v.forward, v.dx))
op.apply(time=6)
exprs = FindNodes(Expression).visit(op)
assert exprs[-1].expr.lhs.indices[0] == IntDiv(time, 2) + 1
def test_symbolics():
a = Symbol('a')
id = IntDiv(a, 3)
pkl_id = pickle.dumps(id)
new_id = pickle.loads(pkl_id)
assert id == new_id
ffp = FunctionFromPointer('foo', a, ['b', 'c'])
pkl_ffp = pickle.dumps(ffp)
new_ffp = pickle.loads(pkl_ffp)
assert ffp == new_ffp
li = ListInitializer(['a', 'b'])
pkl_li = pickle.dumps(li)
new_li = pickle.loads(pkl_li)
assert li == new_li
def _lower_conditional_dims(iet):
"""
Lower ConditionalDimensions: index functions involving ConditionalDimensions
are turned into integer-division expressions.
Examples
--------
u[t_sub, x] = u[time, x]
becomes
u[time / 4, x] = u[time, x]
"""
cdims = [d for d in FindSymbols('free-symbols').visit(iet)
if isinstance(d, ConditionalDimension)]
mapper = {d: IntDiv(d.index, d.factor) for d in cdims}
iet = XSubs(mapper).visit(iet)
return iet
def test_symbolics():
a = Symbol('a')
id = IntDiv(a, 3)
pkl_id = pickle.dumps(id)
new_id = pickle.loads(pkl_id)
assert id == new_id
ffp = CallFromPointer('foo', a, ['b', 'c'])
pkl_ffp = pickle.dumps(ffp)
new_ffp = pickle.loads(pkl_ffp)
assert ffp == new_ffp
li = ListInitializer(['a', 'b'])
pkl_li = pickle.dumps(li)
new_li = pickle.loads(pkl_li)
assert li == new_li
df = DefFunction('f', ['a', 1, 2])
pkl_df = pickle.dumps(df)
new_df = pickle.loads(pkl_df)
assert df == new_df
assert df.arguments == new_df.arguments
def __new__(cls, *args, **kwargs):
if len(args) == 1 and isinstance(args[0], LoweredEq):
# origin: LoweredEq(devito.LoweredEq, **kwargs)
input_expr = args[0]
expr = sympy.Eq.__new__(cls, *input_expr.args, evaluate=False)
for i in cls._state:
setattr(expr, '_%s' % i,
kwargs.get(i) or getattr(input_expr, i))
return expr
elif len(args) == 1 and isinstance(args[0], Eq):
# origin: LoweredEq(devito.Eq)
input_expr = expr = args[0]
elif len(args) == 2:
expr = sympy.Eq.__new__(cls, *args, evaluate=False)
for i in cls._state:
setattr(expr, '_%s' % i, kwargs.pop(i))
return expr
else:
raise ValueError("Cannot construct LoweredEq from args=%s "
"and kwargs=%s" % (str(args), str(kwargs)))
# Well-defined dimension ordering
ordering = dimension_sort(expr)
# Analyze the expression
accesses = detect_accesses(expr)
dimensions = Stencil.union(*accesses.values())
# Separate out the SubIterators from the main iteration Dimensions, that
# is those which define an actual iteration space
iterators = {}
for d in dimensions:
if d.is_SubIterator:
iterators.setdefault(d.root, set()).add(d)
elif d.is_Conditional:
# Use `parent`, and `root`, because a ConditionalDimension may
# have a SubDimension as parent
iterators.setdefault(d.parent, set())
else:
iterators.setdefault(d, set())
# Construct the IterationSpace
intervals = IntervalGroup([Interval(d, 0, 0) for d in iterators],
relations=ordering.relations)
ispace = IterationSpace(intervals, iterators)
# Construct the conditionals and replace the ConditionalDimensions in `expr`
conditionals = {}
for d in ordering:
if not d.is_Conditional:
continue
if d.condition is None:
conditionals[d] = GuardFactor(d)
else:
conditionals[d] = diff2sympy(lower_exprs(d.condition))
if d.factor is not None:
expr = uxreplace(expr, {d: IntDiv(d.index, d.factor)})
conditionals = frozendict(conditionals)
# Lower all Differentiable operations into SymPy operations
rhs = diff2sympy(expr.rhs)
# Finally create the LoweredEq with all metadata attached
expr = super(LoweredEq, cls).__new__(cls,
expr.lhs,
rhs,
evaluate=False)
expr._ispace = ispace
expr._conditionals = conditionals
expr._reads, expr._writes = detect_io(expr)
expr._is_Increment = input_expr.is_Increment
expr._implicit_dims = input_expr.implicit_dims
return expr
def __new__(cls, *args, **kwargs):
if len(args) == 1 and isinstance(args[0], LoweredEq):
# origin: LoweredEq(devito.LoweredEq, **kwargs)
input_expr = args[0]
expr = sympy.Eq.__new__(cls, *input_expr.args, evaluate=False)
for i in cls._state:
setattr(expr, '_%s' % i, kwargs.get(i) or getattr(input_expr, i))
return expr
elif len(args) == 1 and isinstance(args[0], Eq):
# origin: LoweredEq(devito.Eq)
input_expr = expr = args[0]
elif len(args) == 2:
expr = sympy.Eq.__new__(cls, *args, evaluate=False)
for i in cls._state:
setattr(expr, '_%s' % i, kwargs.pop(i))
return expr
else:
raise ValueError("Cannot construct LoweredEq from args=%s "
"and kwargs=%s" % (str(args), str(kwargs)))
# Well-defined dimension ordering
ordering = dimension_sort(expr)
# Analyze the expression
mapper = detect_accesses(expr)
oobs = detect_oobs(mapper)
conditional_dimensions = [i for i in ordering if i.is_Conditional]
# Construct Intervals for IterationSpace and DataSpace
intervals = build_intervals(Stencil.union(*mapper.values()))
iintervals = [] # iteration Intervals
dintervals = [] # data Intervals
for i in intervals:
d = i.dim
if d in oobs:
iintervals.append(i.zero())
dintervals.append(i)
else:
iintervals.append(i.zero())
dintervals.append(i.zero())
# Construct the IterationSpace
iintervals = IntervalGroup(iintervals, relations=ordering.relations)
iterators = build_iterators(mapper)
ispace = IterationSpace(iintervals, iterators)
# Construct the DataSpace
dintervals.extend([Interval(i, 0, 0) for i in ordering
if i not in ispace.dimensions + conditional_dimensions])
parts = {k: IntervalGroup(build_intervals(v)).add(iintervals)
for k, v in mapper.items() if k}
dspace = DataSpace(dintervals, parts)
# Construct the conditionals and replace the ConditionalDimensions in `expr`
conditionals = {}
for d in conditional_dimensions:
if d.condition is None:
conditionals[d] = CondEq(d.parent % d.factor, 0)
else:
conditionals[d] = diff2sympy(lower_exprs(d.condition))
if d.factor is not None:
expr = uxreplace(expr, {d: IntDiv(d.index, d.factor)})
conditionals = frozendict(conditionals)
# Lower all Differentiable operations into SymPy operations
rhs = diff2sympy(expr.rhs)
# Finally create the LoweredEq with all metadata attached
expr = super(LoweredEq, cls).__new__(cls, expr.lhs, rhs, evaluate=False)
expr._dspace = dspace
expr._ispace = ispace
expr._conditionals = conditionals
expr._reads, expr._writes = detect_io(expr)
expr._is_Increment = input_expr.is_Increment
expr._implicit_dims = input_expr.implicit_dims
return expr