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