def clusterize(exprs):
"""
Group a sequence of :class:`ir.Eq`s into one or more :class:`Cluster`s.
"""
clusters = ClusterGroup()
flowmap = detect_flow_directions(exprs)
prev = None
for idx, e in enumerate(exprs):
if e.is_Tensor:
scalars = [i for i in exprs[prev:idx] if i.is_Scalar]
# Iteration space
ispace = IterationSpace.merge(e.ispace, *[i.ispace for i in scalars])
# Enforce iteration directions
fdirs, _ = force_directions(flowmap, lambda d: ispace.directions.get(d))
ispace = IterationSpace(ispace.intervals, ispace.sub_iterators, fdirs)
# Data space
dspace = DataSpace.merge(e.dspace, *[i.dspace for i in scalars])
# Prepare for next range
prev = idx
clusters.append(PartialCluster(scalars + [e], ispace, dspace))
# Group PartialClusters together where possible
clusters = groupby(clusters)
# Introduce conditional PartialClusters
clusters = guard(clusters)
return clusters.finalize()
def squash(self, other):
"""Concatenate the expressions in ``other`` to those in ``self``.
``self`` and ``other`` must have same ``ispace``. Duplicate
expressions are dropped. The :class:`DataSpace` is updated
accordingly."""
assert self.ispace.is_compatible(other.ispace)
self.exprs.extend([i for i in other.exprs if i not in self.exprs])
self.dspace = DataSpace.merge(self.dspace, other.dspace)
self.ispace = IterationSpace.merge(self.ispace, other.ispace)
def squash(self, other):
"""
Concatenate the expressions in ``other`` to those in ``self``.
``self`` and ``other`` must have same ``ispace``. Duplicate expressions
are dropped. The DataSpace is updated accordingly.
"""
assert self.ispace.is_compatible(other.ispace)
self.exprs.extend([i for i in other.exprs
if i not in self.exprs or i.is_Increment])
self.dspace = DataSpace.merge(self.dspace, other.dspace)
self.ispace = IterationSpace.merge(self.ispace, other.ispace)
def from_clusters(cls, *clusters):
"""
Build a new Cluster from a sequence of pre-existing Clusters with
compatible IterationSpace.
"""
assert len(clusters) > 0
root = clusters[0]
assert all(root.ispace.is_compatible(c.ispace) for c in clusters)
exprs = chain(*[c.exprs for c in clusters])
ispace = IterationSpace.merge(*[c.ispace for c in clusters])
dspace = DataSpace.merge(*[c.dspace for c in clusters])
return Cluster(exprs, ispace, dspace)
def st_schedule(clusters):
"""
Arrange an iterable of :class:`Cluster`s into a :class:`ScheduleTree`.
"""
stree = ScheduleTree()
mapper = OrderedDict()
for c in clusters:
pointers = list(mapper)
# Find out if any of the existing nodes can be reused
index = 0
root = stree
for it0, it1 in zip(c.itintervals, pointers):
if it0 != it1 or it0.dim in c.atomics:
break
root = mapper[it0]
index += 1
if it0.dim in c.guards:
break
# The reused sub-trees might acquire some new sub-iterators
for i in pointers[:index]:
mapper[i].ispace = IterationSpace.merge(mapper[i].ispace,
c.ispace.project([i.dim]))
# Later sub-trees, instead, will not be used anymore
for i in pointers[index:]:
mapper.pop(i)
# Add in Iterations
for i in c.itintervals[index:]:
root = NodeIteration(c.ispace.project([i.dim]), root)
mapper[i] = root
# Add in Expressions
NodeExprs(c.exprs, c.ispace, c.dspace, c.shape, c.ops, c.traffic, root)
# Add in Conditionals
for k, v in mapper.items():
if k.dim in c.guards:
node = NodeConditional(c.guards[k.dim])
v.last.parent = node
node.parent = v
return stree
def st_schedule(clusters):
"""
Arrange an iterable of Clusters into a ScheduleTree.
"""
stree = ScheduleTree()
mapper = OrderedDict()
for c in clusters:
pointers = list(mapper)
# Find out if any of the existing nodes can be reused
index = 0
root = stree
for it0, it1 in zip(c.itintervals, pointers):
if it0 != it1 or it0.dim in c.atomics:
break
root = mapper[it0]
index += 1
if it0.dim in c.guards:
break
# The reused sub-trees might acquire some new sub-iterators
for i in pointers[:index]:
mapper[i].ispace = IterationSpace.merge(mapper[i].ispace,
c.ispace.project([i.dim]))
# Later sub-trees, instead, will not be used anymore
for i in pointers[index:]:
mapper.pop(i)
# Add in Iterations
for i in c.itintervals[index:]:
root = NodeIteration(c.ispace.project([i.dim]), root)
mapper[i] = root
# Add in Expressions
NodeExprs(c.exprs, c.ispace, c.dspace, c.shape, c.ops, c.traffic, root)
# Add in Conditionals
for k, v in mapper.items():
if k.dim in c.guards:
node = NodeConditional(c.guards[k.dim])
v.last.parent = node
node.parent = v
return stree
def iet_make(clusters):
"""
Create an Iteration/Expression tree (IET) given an iterable of :class:`Cluster`s.
:param clusters: The iterable :class:`Cluster`s for which the IET is built.
"""
# {Iteration -> [c0, c1, ...]}, shared clusters
shared = {}
# The constructed IET
processed = []
# {Interval -> Iteration}, carried from preceding cluster
schedule = OrderedDict()
# Build IET
for cluster in clusters:
body = [Expression(e) for e in cluster.exprs]
if cluster.ispace.empty:
# No Iterations are needed
processed.extend(body)
continue
root = None
itintervals = cluster.ispace.iteration_intervals
# Can I reuse any of the previously scheduled Iterations ?
index = 0
for i0, i1 in zip(itintervals, list(schedule)):
if i0 != i1 or i0.dim in cluster.atomics:
break
root = schedule[i1]
index += 1
needed = itintervals[index:]
# Build Expressions
if not needed:
body = List(body=body)
# Build Iterations
scheduling = []
for i in reversed(needed):
# Update IET and scheduling
if i.dim in cluster.guards:
# Must wrap within an if-then scope
body = Conditional(cluster.guards[i.dim], body)
# Adding (None, None) ensures that nested iterations won't
# be reused by the next cluster
scheduling.insert(0, (None, None))
iteration = Iteration(body,
i.dim,
i.dim.limits,
offsets=i.limits,
direction=i.direction)
scheduling.insert(0, (i, iteration))
# Prepare for next dimension
body = iteration
# If /needed/ is != [], root.dim might be a guarded dimension for /cluster/
if root is not None and root.dim in cluster.guards:
body = Conditional(cluster.guards[root.dim], body)
# Update the current schedule
if root is None:
processed.append(body)
else:
nodes = list(root.nodes) + [body]
transf = Transformer(
{root: root._rebuild(nodes, **root.args_frozen)})
processed = list(transf.visit(processed))
scheduling = list(schedule.items())[:index] + list(scheduling)
scheduling = [(k, transf.rebuilt.get(v, v)) for k, v in scheduling]
shared = {transf.rebuilt.get(k, k): v for k, v in shared.items()}
schedule = OrderedDict(scheduling)
# Record that /cluster/ was used to build the iterations in /schedule/
shared.update(
{i: shared.get(i, []) + [cluster]
for i in schedule.values() if i})
iet = List(body=processed)
# Add in unbounded indices, if needed
mapper = {}
for k, v in shared.items():
uindices = []
ispace = IterationSpace.merge(*[i.ispace.project([k.dim]) for i in v])
for j, offs in ispace.sub_iterators.get(k.dim, []):
modulo = len(offs)
for n, o in enumerate(filter_ordered(offs)):
name = "%s%d" % (j.name, n)
vname = Scalar(name=name, dtype=np.int32)
value = (k.dim + o) % modulo
uindices.append(UnboundedIndex(vname, value, value, j, j + o))
mapper[k] = k._rebuild(uindices=uindices)
iet = NestedTransformer(mapper).visit(iet)
return iet