def clusterize(exprs, stencils):
"""Derive :class:`Cluster`s from an iterator of expressions; a stencil for
each expression must be provided."""
assert len(exprs) == len(stencils)
exprs, stencils = aggregate(exprs, stencils)
g = temporaries_graph(exprs)
mapper = OrderedDict([(i.lhs, j) for i, j in zip(g.values(), stencils)
if i.is_tensor])
clusters = []
for k, v in mapper.items():
# Determine what temporaries are needed to compute /i/
exprs = g.trace(k)
# Determine the Stencil of the cluster
stencil = Stencil(v.entries)
for i in exprs:
stencil = stencil.add(mapper.get(i.lhs, {}))
stencil = stencil.frozen
# Drop all non-output tensors, as computed by other clusters
exprs = [i for i in exprs if i.lhs.is_Symbol or i.lhs == k]
# Create and track the cluster
clusters.append(Cluster(exprs, stencil))
return merge(clusters)
def clusterize(exprs, stencils, atomics=None):
"""
Derive :class:`Cluster` objetcs from an iterator of expressions; a stencil for
each expression must be provided. A list of atomic dimensions (see description
in Cluster.__doc__) may be provided.
"""
assert len(exprs) == len(stencils)
exprs, stencils = aggregate(exprs, stencils)
Info = namedtuple('Info', 'trace stencil')
g = temporaries_graph(exprs)
mapper = OrderedDict([
(k, Info(g.trace(k) + g.trace(k, readby=True, strict=True), j))
for (k, v), j in zip(g.items(), stencils) if v.is_tensor
])
# A cluster stencil is determined iteratively, by first calculating the
# "local" stencil and then by looking at the stencils of all other clusters
# depending on it. The stencil information is propagated until there are
# no more updates.
queue = list(mapper)
while queue:
target = queue.pop(0)
info = mapper[target]
strict_trace = [i.lhs for i in info.trace if i.lhs != target]
stencil = Stencil(info.stencil.entries)
for i in strict_trace:
if i in mapper:
stencil = stencil.add(mapper[i].stencil)
mapper[target] = Info(info.trace, stencil)
if stencil != info.stencil:
# Something has changed, need to propagate the update
queue.extend([i for i in strict_trace if i not in queue])
clusters = []
for target, info in mapper.items():
# Drop all non-output tensors, as computed by other clusters
exprs = [i for i in info.trace if i.lhs.is_Symbol or i.lhs == target]
# Create and track the cluster
clusters.append(Cluster(exprs, info.stencil.frozen, atomics))
return merge(clusters)
def anti_stencil(self):
handle = Stencil()
for d, i in zip(self.dimensions, zip(*self.distances)):
handle[d].update(set(i))
for d, i in zip(self.dimensions, zip(*self._ghost_offsets)):
handle[d].update(set(i))
return handle
def _retrieve_loop_ordering(self, expressions):
"""
Establish a partial ordering for the loops that will appear in the code
generated by the Operator, based on the order in which dimensions
appear in the input expressions.
"""
ordering = []
for i in flatten(Stencil(i).dimensions for i in expressions):
if i not in ordering:
ordering.extend([i, i.parent] if i.is_Buffered else [i])
return ordering
def _retrieve_stencils(self, expressions):
"""Determine the :class:`Stencil` of each provided expression."""
stencils = [Stencil(i) for i in expressions]
dimensions = set.union(*[set(i.dimensions) for i in stencils])
# Filter out aliasing buffered dimensions
mapper = {d.parent: d for d in dimensions if d.is_Buffered}
for i in list(stencils):
for d in i.dimensions:
if d in mapper:
i[mapper[d]] = i.pop(d).union(i.get(mapper[d], set()))
return stencils
def merge(clusters):
"""
Given an ordered collection of :class:`Cluster` objects, return a
(potentially) smaller sequence in which clusters with identical stencil
have been merged into a single :class:`Cluster`.
"""
mapper = OrderedDict()
for c in clusters:
mapper.setdefault(c.stencil.entries, []).append(c)
processed = []
for entries, clusters in mapper.items():
# Eliminate redundant temporaries
temporaries = OrderedDict()
for c in clusters:
for k, v in c.trace.items():
if k not in temporaries:
temporaries[k] = v
# Squash the clusters together
processed.append(Cluster(temporaries.values(), Stencil(entries)))
return processed
def stencil(self):
"""Compute the stencil of the expression."""
return Stencil(self.expr)
def collect_aliases(exprs):
"""
Determine groups of aliasing expressions in ``exprs``.
An expression A aliases an expression B if both A and B apply the same
operations to the same input operands, with the possibility for indexed objects
to index into locations at a fixed constant offset in each dimension.
For example: ::
exprs = (a[i+1] + b[i+1], a[i+1] + b[j+1], a[i] + c[i],
a[i+2] - b[i+2], a[i+2] + b[i], a[i-1] + b[i-1])
The following expressions in ``exprs`` alias to ``a[i] + b[i]``: ::
a[i+1] + b[i+1] : same operands and operations, distance along i = 1
a[i-1] + b[i-1] : same operands and operations, distance along i = -1
Whereas the following do not: ::
a[i+1] + b[j+1] : because at least one index differs
a[i] + c[i] : because at least one of the operands differs
a[i+2] - b[i+2] : because at least one operation differs
a[i+2] + b[i] : because distance along ``i`` differ (+2 and +0)
"""
ExprData = namedtuple('ExprData', 'dimensions offsets')
# Discard expressions that surely won't alias to anything
candidates = OrderedDict()
for expr in exprs:
indexeds = retrieve_indexed(expr, mode='all')
if indexeds and not any(q_indirect(i) for i in indexeds):
handle = calculate_offsets(indexeds)
if handle:
candidates[expr.rhs] = ExprData(*handle)
aliases = OrderedDict()
mapper = OrderedDict()
unseen = list(candidates)
while unseen:
# Find aliasing expressions
handle = unseen.pop(0)
group = [handle]
for e in list(unseen):
if compare(handle, e) and\
is_translated(candidates[handle].offsets, candidates[e].offsets):
group.append(e)
unseen.remove(e)
mapper.update([(i, group) for i in group])
if len(group) == 1:
continue
# Create an alias for the group of aliasing expressions, as well as
# any metadata needed by the caller
offsets = [tuple(candidates[e].offsets) for e in group]
COM, distances = calculate_COM(offsets)
alias = create_alias(handle, COM)
aliases[alias] = Alias(alias, group, distances, candidates[handle].dimensions)
# Heuristically attempt to relax the aliases offsets
# to maximize the likelyhood of loop fusion
grouped = OrderedDict()
for i in aliases.values():
grouped.setdefault(i.dimensions, []).append(i)
for dimensions, group in grouped.items():
ideal_anti_stencil = Stencil.union(*[i.anti_stencil for i in group])
for i in group:
if i.anti_stencil.subtract(ideal_anti_stencil).empty:
aliases[i.alias] = i.relax(ideal_anti_stencil)
return mapper, aliases
def collect(exprs):
"""
Determine groups of aliasing expressions in ``exprs``.
An expression A aliases an expression B if both A and B apply the same
operations to the same input operands, with the possibility for indexed objects
to index into locations at a fixed constant offset in each dimension.
For example: ::
exprs = (a[i+1] + b[i+1], a[i+1] + b[j+1], a[i] + c[i],
a[i+2] - b[i+2], a[i+2] + b[i], a[i-1] + b[i-1])
The following expressions in ``exprs`` alias to ``a[i] + b[i]``: ::
a[i+1] + b[i+1] : same operands and operations, distance along i = 1
a[i-1] + b[i-1] : same operands and operations, distance along i = -1
Whereas the following do not: ::
a[i+1] + b[j+1] : because at least one index differs
a[i] + c[i] : because at least one of the operands differs
a[i+2] - b[i+2] : because at least one operation differs
a[i+2] + b[i] : because distance along ``i`` differ (+2 and +0)
"""
ExprData = namedtuple('ExprData', 'dimensions offsets')
# Discard expressions:
# - that surely won't alias to anything
# - that are non-scalar
candidates = OrderedDict()
for expr in exprs:
if q_indexed(expr):
continue
indexeds = retrieve_indexed(expr.rhs, mode='all')
if indexeds and not any(q_indirect(i) for i in indexeds):
handle = calculate_offsets(indexeds)
if handle:
candidates[expr.rhs] = ExprData(*handle)
aliases = OrderedDict()
mapper = OrderedDict()
unseen = list(candidates)
while unseen:
# Find aliasing expressions
handle = unseen.pop(0)
group = [handle]
for e in list(unseen):
if compare(handle, e) and\
is_translated(candidates[handle].offsets, candidates[e].offsets):
group.append(e)
unseen.remove(e)
mapper.update([(i, group) for i in group])
# Try creating a basis for the aliasing expressions' offsets
offsets = [tuple(candidates[e].offsets) for e in group]
try:
COM, distances = calculate_COM(offsets)
except DSEException:
# Ignore these potential aliases and move on
continue
alias = create_alias(handle, COM)
# In circumstances in which an expression has repeated coefficients, e.g.
# ... + 0.025*a[...] + 0.025*b[...],
# We may have found a common basis (i.e., same COM, same alias) at this point
v = aliases.setdefault(alias, Alias(alias, candidates[handle].dimensions))
v.extend(group, distances)
# Heuristically attempt to relax the aliases offsets
# to maximize the likelyhood of loop fusion
grouped = OrderedDict()
for i in aliases.values():
grouped.setdefault(i.dimensions, []).append(i)
for dimensions, group in grouped.items():
ideal_anti_stencil = Stencil.union(*[i.anti_stencil for i in group])
for i in group:
if i.anti_stencil.subtract(ideal_anti_stencil).empty:
aliases[i.alias] = i.relax(ideal_anti_stencil)
return mapper, aliases
@pytest.mark.parametrize(
'exprs,expected',
[
# none (different distance)
(['Eq(t0, fa[x] + fb[x])', 'Eq(t1, fa[x+1] + fb[x])'], {}),
# none (different dimension)
(['Eq(t0, fa[x] + fb[x])', 'Eq(t1, fa[x] + fb[y])'], {}),
# none (different operation)
(['Eq(t0, fa[x] + fb[x])', 'Eq(t1, fa[x] - fb[x])'], {}),
# simple
([
'Eq(t0, fa[x] + fb[x])', 'Eq(t1, fa[x+1] + fb[x+1])',
'Eq(t2, fa[x-1] + fb[x-1])'
], {
'fa[x] + fb[x]': Stencil([(x, {-1, 0, 1})])
}),
# 2D simple
(['Eq(t0, fc[x,y] + fd[x,y])', 'Eq(t1, fc[x+1,y+1] + fd[x+1,y+1])'], {
'fc[x,y] + fd[x,y]': Stencil([(x, {0, 1}), (y, {0, 1})])
}),
# 2D with stride
([
'Eq(t0, fc[x,y] + fd[x+1,y+2])',
'Eq(t1, fc[x+1,y+1] + fd[x+2,y+3])'
], {
'fc[x,y] + fd[x+1,y+2]': Stencil([(x, {0, 1}), (y, {0, 1})])
}),
# complex (two 2D aliases with stride inducing relaxation)
([
'Eq(t0, fc[x,y] + fd[x+1,y+2])',