def _lower_exprs(cls, expressions, **kwargs):
"""
Expression lowering:
* Form and gather any required implicit expressions;
* Evaluate derivatives;
* Flatten vectorial equations;
* Indexify Functions;
* Apply substitution rules;
* Specialize (e.g., index shifting)
"""
# Add in implicit expressions, e.g., induced by SubDomains
expressions = cls._add_implicit(expressions)
# Unfold lazyiness
expressions = flatten([i.evaluate for i in expressions])
# Scalarize tensor expressions
expressions = [j for i in expressions for j in i._flatten]
expressions = lower_exprs(expressions, **kwargs)
processed = cls._specialize_exprs(expressions)
return processed
def guard(clusters):
"""
Split Clusters containing conditional expressions into separate Clusters.
"""
processed = []
for c in clusters:
# Group together consecutive expressions with same ConditionalDimensions
for cds, g in groupby(c.exprs, key=lambda e: e.conditionals):
if not cds:
processed.append(c.rebuild(exprs=list(g)))
continue
# Create a guarded Cluster
guards = {}
for cd in cds:
condition = guards.setdefault(cd.parent, [])
if cd.condition is None:
condition.append(CondEq(cd.parent % cd.factor, 0))
else:
condition.append(lower_exprs(cd.condition))
guards = {
k: sympy.And(*v, evaluate=False)
for k, v in guards.items()
}
processed.append(c.rebuild(exprs=list(g), guards=guards))
return ClusterGroup(processed)
def _lower_exprs(cls, expressions, **kwargs):
"""
Expression lowering:
* Form and gather any required implicit expressions;
* Apply rewrite rules;
* Evaluate derivatives;
* Flatten vectorial equations;
* Indexify Functions;
* Apply substitution rules;
* Shift indices for domain alignment.
"""
# Add in implicit expressions
expressions = generate_implicit_exprs(expressions)
# Specialization is performed on unevaluated expressions
expressions = cls._specialize_exprs(expressions, **kwargs)
# Lower functional DSL
expressions = flatten([i.evaluate for i in expressions])
expressions = [j for i in expressions for j in i._flatten]
# "True" lowering (indexification, shifting, ...)
expressions = lower_exprs(expressions, **kwargs)
processed = [LoweredEq(i) for i in expressions]
return processed
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
conditionals = {}
for d in conditional_dimensions:
if d.condition is None:
conditionals[d] = CondEq(d.parent % d.factor, 0)
else:
conditionals[d] = lower_exprs(d.condition)
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
